test shift_operator.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bbd35cc9-9e73-4906-b9e4-3ee83cf2aae2",
   "metadata": {},
   "source": [
    "# Test for `shift_operator` of jwave / jaxdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "99ac4b34-b57b-425f-9069-09c20be261ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax.numpy as jnp\n",
    "from jwave import Domain, FiniteDifferences, FourierSeries, Continuous"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1252d2a3-bd0c-445b-adc4-1a2b5a263a47",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jaxdf\n",
    "from jaxdf.operators.functions import shift_operator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ca0b1fd5-e432-468d-9d71-dd7da5f2dfe7",
   "metadata": {},
   "outputs": [],
   "source": [
    "jnp.set_printoptions(precision=3, floatmode='maxprec', suppress=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7993b6cf-f3fe-4723-aaff-b1ff1463e147",
   "metadata": {},
   "source": [
    "## initialize field\n",
    "\n",
    "initialize fields: FourierSeries and FiniteDifferences for scalar field (single component) and vector field (2 components)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "850449ec-1647-46d2-8990-b069f1148565",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = jnp.zeros((5,5))\n",
    "a = a.at[2,2].set(100)\n",
    "a_s = a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d1e68256-66ff-4887-98da-cc76d3f28524",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = jnp.zeros((3,3,2))\n",
    "a = a.at[1,1,0].set(10)\n",
    "a = a.at[1,1,1].set(20)\n",
    "a_v = a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "97e153f2-c321-47d3-9fa4-deadf25468ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "domain_s = Domain(a_s.shape[:2], dx = [1,1])\n",
    "domain_v = Domain(a_v.shape[:2], dx = [1,1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7dadb598-6ad4-4461-acbc-dea3dec8feb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "fF_s = FourierSeries(a_s, domain_s)\n",
    "fF_v = FourierSeries(a_v, domain_v)\n",
    "fD_s = FiniteDifferences(a_s, domain_s)\n",
    "fD_v = FiniteDifferences(a_v, domain_v)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42caca79-43dd-41d1-a0a0-cc8a9299edec",
   "metadata": {},
   "source": [
    "pretty printing field components individually"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f8f26198-fb60-495c-aee2-e13645c87f9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pretty printing field components\n",
    "def pp(field):\n",
    "    print(jnp.rollaxis(field.on_grid, -1, 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "59a3ed7c-22e2-42a9-8e35-f7d3220a8071",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[  0.   0.   0.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]\n",
      "  [  0.   0. 100.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]]]\n"
     ]
    }
   ],
   "source": [
    "pp(fF_s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "883edc00-9547-4599-a802-b7ba4c740f5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[ 0.  0.  0.]\n",
      "  [ 0. 10.  0.]\n",
      "  [ 0.  0.  0.]]\n",
      "\n",
      " [[ 0.  0.  0.]\n",
      "  [ 0. 20.  0.]\n",
      "  [ 0.  0.  0.]]]\n"
     ]
    }
   ],
   "source": [
    "pp(fF_v)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca1d16a9-31ee-454f-a7f2-f0939283f2f7",
   "metadata": {},
   "source": [
    "## `dx = [1]`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "1e45f28f-8f4f-4971-8c43-00adfe292d07",
   "metadata": {},
   "outputs": [],
   "source": [
    "dx = [1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6e5a6b3d-851c-4bc0-88b4-706d5db5e92c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[ 0. 10.  0.]\n",
      "  [ 0.  0.  0.]\n",
      "  [ 0.  0.  0.]]\n",
      "\n",
      " [[ 0.  0.  0.]\n",
      "  [ 0.  0. 20.]\n",
      "  [ 0.  0.  0.]]]\n"
     ]
    }
   ],
   "source": [
    "# FiniteDifferences\n",
    "pp(shift_operator(fD_v, dx=dx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "28a1ec13-4146-4c80-955d-31c30c46aa71",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-08-01 12:06:30 - jaxdf [WARNING]: Deprecation: Currently only the first output of an operator is considered. This will change in a future release. If you need to return multiple outputs, please return a tuple and a None value, for example: ((out1, out2), None). This happened for the operator `shift_operator`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[ 0. 10.  0.]\n",
      "  [ 0. -0.  0.]\n",
      "  [ 0.  0.  0.]]\n",
      "\n",
      " [[ 0.  0.  0.]\n",
      "  [20. -0.  0.]\n",
      "  [ 0.  0.  0.]]]\n"
     ]
    }
   ],
   "source": [
    "# FourierSeries\n",
    "pp(shift_operator(fF_v, dx=dx))\n",
    "\n",
    "# acts in different direction!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "906f4e44-4aeb-41b6-904b-749dc99df3ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[  0.   0.   0.   0.   0.]\n",
      "  [  0.   0. 100.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]]\n",
      "\n",
      " [[  0.   0.   0.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]\n",
      "  [  0.   0.   0. 100.   0.]\n",
      "  [  0.   0.   0.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]]]\n"
     ]
    }
   ],
   "source": [
    "# FiniteDifferences, scalar\n",
    "pp(shift_operator(fD_s, dx=dx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1a67261d-a513-42ed-b27e-732b67ea31c0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-08-01 12:06:32 - jaxdf [WARNING]: Deprecation: Currently only the first output of an operator is considered. This will change in a future release. If you need to return multiple outputs, please return a tuple and a None value, for example: ((out1, out2), None). This happened for the operator `shift_operator`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[  0.   0.   0.   0.   0.]\n",
      "  [  0.   0. 100.   0.   0.]\n",
      "  [  0.   0.  -0.   0.   0.]\n",
      "  [  0.   0.  -0.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]]]\n"
     ]
    }
   ],
   "source": [
    "# FourierSeries, scalar\n",
    "pp(shift_operator(fF_s, dx=dx)) # fails"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7d9d15d-dd4f-4baf-833b-734cffc06d1e",
   "metadata": {},
   "source": [
    "## `dx = [1, 0]`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c39cf2ed-2dd6-4aa2-8eb7-042a70e763aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "dx = [1, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "24ef5897-25d1-4998-94d7-bb602e8e9518",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[ 0. 10.  0.]\n",
      "  [ 0.  0.  0.]\n",
      "  [ 0.  0.  0.]]\n",
      "\n",
      " [[ 0.  0.  0.]\n",
      "  [ 0.  0. 20.]\n",
      "  [ 0.  0.  0.]]]\n"
     ]
    }
   ],
   "source": [
    "# FiniteDifferences, vectorial\n",
    "pp(shift_operator(fD_v, dx=dx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "589f478a-7fd7-4355-b3c6-911e6c2c9e7b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-08-01 12:06:49 - jaxdf [WARNING]: Deprecation: Currently only the first output of an operator is considered. This will change in a future release. If you need to return multiple outputs, please return a tuple and a None value, for example: ((out1, out2), None). This happened for the operator `shift_operator`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[ 0. 10.  0.]\n",
      "  [ 0. -0.  0.]\n",
      "  [ 0.  0.  0.]]\n",
      "\n",
      " [[ 0.  0.  0.]\n",
      "  [ 0. 20.  0.]\n",
      "  [ 0.  0.  0.]]]\n"
     ]
    }
   ],
   "source": [
    "# FourierSeries, scalar\n",
    "pp(shift_operator(fF_v, dx=dx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "f088b40c-9804-461f-b788-a15afc112274",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[  0.   0.   0.   0.   0.]\n",
      "  [  0.   0. 100.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]]\n",
      "\n",
      " [[  0.   0.   0.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]\n",
      "  [  0.   0.   0. 100.   0.]\n",
      "  [  0.   0.   0.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]]]\n"
     ]
    }
   ],
   "source": [
    "# FiniteDifferences, scalar\n",
    "pp(shift_operator(fD_s, dx=dx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "4a39646b-b931-4c65-b5f5-6bc38f4e61cc",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-08-01 12:07:04 - jaxdf [WARNING]: Deprecation: Currently only the first output of an operator is considered. This will change in a future release. If you need to return multiple outputs, please return a tuple and a None value, for example: ((out1, out2), None). This happened for the operator `shift_operator`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[  0.   0.   0.   0.   0.]\n",
      "  [  0.   0. 100.   0.   0.]\n",
      "  [  0.   0.  -0.   0.   0.]\n",
      "  [  0.   0.  -0.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]]]\n"
     ]
    }
   ],
   "source": [
    "# FourierSeries, scalar\n",
    "pp(shift_operator(fF_s, dx=dx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbe8be3a-6e86-4aeb-8664-aad678866890",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "759a1d93-2e54-46eb-959b-6f12357713dc",
   "metadata": {},
   "source": [
    "## `dx = [0, 1]`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "873efae1-7628-43a5-a034-7a697bd3d01e",
   "metadata": {},
   "outputs": [],
   "source": [
    "dx = [0, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "35b8e0b0-692e-47ae-87a8-df2dddde572b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[ 0.  0.  0.]\n",
      "  [ 0. 10.  0.]\n",
      "  [ 0.  0.  0.]]\n",
      "\n",
      " [[ 0.  0.  0.]\n",
      "  [ 0. 20.  0.]\n",
      "  [ 0.  0.  0.]]]\n"
     ]
    }
   ],
   "source": [
    "# FiniteDifferences, vectorial\n",
    "pp(shift_operator(fD_v, dx=dx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "7fc2c90b-8bb9-4e99-839c-13f05d395d38",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-08-01 12:07:07 - jaxdf [WARNING]: Deprecation: Currently only the first output of an operator is considered. This will change in a future release. If you need to return multiple outputs, please return a tuple and a None value, for example: ((out1, out2), None). This happened for the operator `shift_operator`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[ 0.  0.  0.]\n",
      "  [ 0. 10.  0.]\n",
      "  [ 0.  0.  0.]]\n",
      "\n",
      " [[ 0.  0.  0.]\n",
      "  [20. -0.  0.]\n",
      "  [ 0.  0.  0.]]]\n"
     ]
    }
   ],
   "source": [
    "# FourierSeries, scalar\n",
    "pp(shift_operator(fF_v, dx=dx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "e346cba3-51d0-43a6-8454-430cc9c45957",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[  0.   0.   0.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]\n",
      "  [  0.   0. 100.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]]\n",
      "\n",
      " [[  0.   0.   0.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]\n",
      "  [  0.   0. 100.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]]]\n"
     ]
    }
   ],
   "source": [
    "# FiniteDifferences, scalar\n",
    "pp(shift_operator(fD_s, dx=dx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "e19c64e5-f716-42e3-a543-40a9b4fb1270",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-08-01 12:07:09 - jaxdf [WARNING]: Deprecation: Currently only the first output of an operator is considered. This will change in a future release. If you need to return multiple outputs, please return a tuple and a None value, for example: ((out1, out2), None). This happened for the operator `shift_operator`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[[  0.   0.  -0.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]\n",
      "  [  0.   0. 100.   0.   0.]\n",
      "  [  0.   0.   0.   0.   0.]\n",
      "  [  0.   0.  -0.   0.   0.]]]\n"
     ]
    }
   ],
   "source": [
    "# FourierSeries, scalar\n",
    "pp(shift_operator(fF_s, dx=dx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f56b537-fd84-46cd-8bf5-3df7478095ac",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (3.13)",
   "language": "python",
   "name": "py313"
  },
  "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}