{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ae32f312",
   "metadata": {},
   "source": [
    "# Conditional logic as array operations – `where`\n",
    "\n",
    "The  [numpy.where](https://numpy.org/doc/stable/reference/generated/numpy.where.html) function is a vectorised version of `if` and `else`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86aefffc",
   "metadata": {},
   "source": [
    "In the following example, we first create a Boolean array and two arrays with values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "94a0f898",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-15T17:40:37.260301Z",
     "iopub.status.busy": "2026-05-15T17:40:37.260050Z",
     "iopub.status.idle": "2026-05-15T17:40:37.295968Z",
     "shell.execute_reply": "2026-05-15T17:40:37.295536Z",
     "shell.execute_reply.started": "2026-05-15T17:40:37.260220Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d97f45df-7602-4967-86d9-6994b1d9570a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-15T17:40:37.296899Z",
     "iopub.status.busy": "2026-05-15T17:40:37.296704Z",
     "iopub.status.idle": "2026-05-15T17:40:37.305948Z",
     "shell.execute_reply": "2026-05-15T17:40:37.305596Z",
     "shell.execute_reply.started": "2026-05-15T17:40:37.296888Z"
    }
   },
   "outputs": [],
   "source": [
    "cond = [False, True, False, True, False, False, False]\n",
    "default_rng = np.random.default_rng()\n",
    "data1 = default_rng.random(7)\n",
    "data2 = default_rng.random(7)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22e59173",
   "metadata": {},
   "source": [
    "Now we want to take the values from `data1` if the corresponding value in `cond` is `True` and otherwise take the value from `data2`. With Python’s `if-else`, this could look like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f3da08da",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-15T17:40:37.306489Z",
     "iopub.status.busy": "2026-05-15T17:40:37.306376Z",
     "iopub.status.idle": "2026-05-15T17:40:37.310164Z",
     "shell.execute_reply": "2026-05-15T17:40:37.309836Z",
     "shell.execute_reply.started": "2026-05-15T17:40:37.306480Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.88985185, 0.34568582, 0.76929043, 0.67471859, 0.26424192,\n",
       "       0.6908897 , 0.88940803])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result = np.array(\n",
    "    [(x if c else y) for x, y, c in zip(data1, data2, cond, strict=True)]\n",
    ")\n",
    "\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dd95e6b",
   "metadata": {},
   "source": [
    "However, this has the following two problems:\n",
    "\n",
    "* with large arrays the function will not be very fast\n",
    "* this will not work with multidimensional arrays"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7540d25f",
   "metadata": {},
   "source": [
    "With `np.where` you can work around these problems in a single function call:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "62fd7e9d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-15T17:40:37.310667Z",
     "iopub.status.busy": "2026-05-15T17:40:37.310592Z",
     "iopub.status.idle": "2026-05-15T17:40:37.312912Z",
     "shell.execute_reply": "2026-05-15T17:40:37.312572Z",
     "shell.execute_reply.started": "2026-05-15T17:40:37.310661Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.88985185, 0.34568582, 0.76929043, 0.67471859, 0.26424192,\n",
       "       0.6908897 , 0.88940803])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result = np.where(cond, data1, data2)\n",
    "\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e00c5267",
   "metadata": {},
   "source": [
    "The second and third arguments of `np.where` do not have to be arrays; one or both can also be scalars. A typical use of `where` in data analysis is to create a new array of values based on another array. Suppose you have a matrix of randomly generated data and you want to make all the negative values positive values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4fbe852d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-15T17:40:37.313289Z",
     "iopub.status.busy": "2026-05-15T17:40:37.313219Z",
     "iopub.status.idle": "2026-05-15T17:40:37.316355Z",
     "shell.execute_reply": "2026-05-15T17:40:37.316128Z",
     "shell.execute_reply.started": "2026-05-15T17:40:37.313283Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.36584626, 0.91970031, 0.13618506, 0.24056537],\n",
       "       [0.59235974, 0.40813171, 0.71144251, 0.08058474],\n",
       "       [0.21988608, 0.50140418, 0.41486371, 0.8820743 ],\n",
       "       [0.4000475 , 0.85196835, 0.33001728, 0.07373272]])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = default_rng.random(size=(4, 4))\n",
    "\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "aacc3846",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-15T17:40:37.317914Z",
     "iopub.status.busy": "2026-05-15T17:40:37.317720Z",
     "iopub.status.idle": "2026-05-15T17:40:37.320215Z",
     "shell.execute_reply": "2026-05-15T17:40:37.319934Z",
     "shell.execute_reply.started": "2026-05-15T17:40:37.317906Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[False,  True, False, False],\n",
       "       [False, False, False, False],\n",
       "       [False, False, False, False],\n",
       "       [False, False, False, False]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data > 0.9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "429ea702",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-15T17:40:37.320771Z",
     "iopub.status.busy": "2026-05-15T17:40:37.320647Z",
     "iopub.status.idle": "2026-05-15T17:40:37.323045Z",
     "shell.execute_reply": "2026-05-15T17:40:37.322750Z",
     "shell.execute_reply.started": "2026-05-15T17:40:37.320754Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.36584626,        nan, 0.13618506, 0.24056537],\n",
       "       [0.59235974, 0.40813171, 0.71144251, 0.08058474],\n",
       "       [0.21988608, 0.50140418, 0.41486371, 0.8820743 ],\n",
       "       [0.4000475 , 0.85196835, 0.33001728, 0.07373272]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.where(data > 0.9, np.nan, data)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.13 Kernel",
   "language": "python",
   "name": "python313"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.0"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}