Adding sparse support

experiments/batch/2-norm.ipynb

@@ -0,0 +1,105 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "9eeb9d3f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import jaxquantum as jqt\n",
+    "import jax.numpy as jnp "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "1b2f2439",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def test_norm_unit():\n",
+    "\n",
+    "    norm_test_array = jnp.array([[1.0, 2.0, 3.0],[4.0, 5.0, 6.0]])\n",
+    "    ones_test_array = jnp.ones_like(norm_test_array)\n",
+    "\n",
+    "    qarr = jqt.coherent(5, jnp.array([[1.0, 2.0, 3.0],[4.0, 5.0, 6.0]]))\n",
+    "    qarr = qarr / norm_test_array\n",
+    "    \n",
+    "    qarr_dm = qarr.to_dm()\n",
+    "\n",
+    "\n",
+    "    # NORM TESTS ----\n",
+    "\n",
+    "    # Case 1: batched dm/operator\n",
+    "    assert jnp.allclose(qarr_dm.norm(), norm_test_array**(-2), rtol=0, atol=1e-6).all()\n",
+    "\n",
+    "    # Case 2: batched ket\n",
+    "    assert jnp.allclose(qarr.norm(), norm_test_array**(-1), rtol=0, atol=1e-6).all()\n",
+    "\n",
+    "    # Case 3: single dm/operator\n",
+    "    assert jnp.isclose(qarr_dm[0,1].norm(), (1/norm_test_array[0,1])**2, rtol=0, atol=1e-6)\n",
+    "\n",
+    "    # Case 4: single ket\n",
+    "    assert jnp.isclose(qarr[0,1].norm(), 1/norm_test_array[0,1], rtol=0, atol=1e-6)\n",
+    "\n",
+    "    # ----\n",
+    "\n",
+    "\n",
+    "    # UNIT TESTS ----\n",
+    "\n",
+    "    # Case 1: batched dm/operator\n",
+    "    assert jnp.allclose(qarr_dm.unit().norm(), ones_test_array, rtol=0, atol=1e-6).all()\n",
+    "\n",
+    "    # Case 2: batched ket\n",
+    "    assert jnp.allclose(qarr.unit().norm(), ones_test_array, rtol=0, atol=1e-6).all()\n",
+    "\n",
+    "    # Case 3: single dm/operator\n",
+    "    assert jnp.isclose(qarr_dm[0,1].unit().norm(), 1.0, rtol=0, atol=1e-6)\n",
+    "\n",
+    "    # Case 4: single ket\n",
+    "    assert jnp.isclose(qarr[0,1].unit().norm(), 1.0, rtol=0, atol=1e-6)\n",
+    "    # ----"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "5bdba403",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "test_norm_unit()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ebdc1a85",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "jax-framework",
+   "language": "python",
+   "name": "python3"
+  },
+  "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.9.18"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

experiments/main/0-examples.ipynb

@@ -740,7 +740,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "jax-0.2.0-env",
+   "display_name": "jax-framework",
    "language": "python",
    "name": "python3"
   },
@@ -754,7 +754,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.5"
+   "version": "3.9.18"
   }
  },
  "nbformat": 4,