{"id":24446384,"url":"https://github.com/i-a-morozov/model","last_synced_at":"2025-10-06T15:30:33.236Z","repository":{"id":254988521,"uuid":"847902860","full_name":"i-a-morozov/model","owner":"i-a-morozov","description":"Simple accelerator lattice model with composable derivatives in pytorch. Computation of parametrix closed orbit, Twiss parameters and other observables.","archived":false,"fork":false,"pushed_at":"2024-10-18T08:29:56.000Z","size":63202,"stargazers_count":0,"open_issues_count":0,"forks_count":0,"subscribers_count":2,"default_branch":"main","last_synced_at":"2024-10-19T11:20:05.696Z","etag":null,"topics":["accelerator-physics","automatic-differentiation","pytorch"],"latest_commit_sha":null,"homepage":"https://i-a-morozov.github.io/model/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/i-a-morozov.png","metadata":{"files":{"readme":"README.MD","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-08-26T19:04:44.000Z","updated_at":"2024-10-18T08:26:35.000Z","dependencies_parsed_at":"2024-09-18T14:55:35.816Z","dependency_job_id":"bdde7c88-bd64-4be2-8487-be2156973a50","html_url":"https://github.com/i-a-morozov/model","commit_stats":null,"previous_names":["i-a-morozov/model"],"tags_count":1,"template":false,"template_full_name":null,"purl":"pkg:github/i-a-morozov/model","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/i-a-morozov%2Fmodel","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/i-a-morozov%2Fmodel/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/i-a-morozov%2Fmodel/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/i-a-morozov%2Fmodel/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/i-a-morozov","download_url":"https://codeload.github.com/i-a-morozov/model/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/i-a-morozov%2Fmodel/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":278635402,"owners_count":26019767,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","status":"online","status_checked_at":"2025-10-06T02:00:05.630Z","response_time":65,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["accelerator-physics","automatic-differentiation","pytorch"],"created_at":"2025-01-20T23:55:19.846Z","updated_at":"2025-10-06T15:30:31.555Z","avatar_url":"https://github.com/i-a-morozov.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# model, 2024\n\n\u003cp align=\"center\"\u003e\n  \u003cimg width=\"300\" height=\"190\" src=\"https://github.com/i-a-morozov/model/blob/main/docs/pics/logo.png\"\u003e\n\u003c/p\u003e\n\n[![DOI](https://zenodo.org/badge/847902860.svg)](https://doi.org/10.5281/zenodo.13933324)\n\nSimple accelerator lattice model with composable derivatives in PyTorch. \nComputation of parametric closed orbit, Twiss parameters and other observables.\n\n# Install\n\n```\n$ pip install git+https://github.com/i-a-morozov/model.git@main\n```\n\n# Documentation and examples\n\n[![Run In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/i-a-morozov/model/blob/main/docs/source/examples/model.ipynb)\n\n[https://i-a-morozov.github.io/model/](https://i-a-morozov.github.io/model/)\n\n# Demo\n\n\n```python\n# In this DEMO basic functionality used for linear lattice analysis is demonstrated\n```\n\n\n```python\n# Set data type and device\n\nimport torch\n\ndtype = torch.float64\ndevice = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n```\n\n\n```python\n# Use selected data type and device for all elements\n\nfrom model.library.element import Element\n\nElement.dtype = dtype\nElement.device = device\n```\n\n\n```python\n# Load ELEGANT SDDS twiss or other table\n\nfrom pathlib import Path\nfrom model.command.external import load_sdds\n\npath = Path('ic.twiss')\n\n# parameters -- {parameter: value}\n# columns    -- {element: {parameter: value}}\n\nparameters, columns = load_sdds(path)\n```\n\n\n```python\n# Set tunes, twiss parameters, phase advances and dispersion at BPMs\n\nfrom model.command.util import select\n\nNUX = torch.tensor(parameters['nux'] % 1, dtype=torch.float64)\nNUY = torch.tensor(parameters['nuy'] % 1, dtype=torch.float64)\n\nkinds = select(columns, 'ElementType')\n\nS = torch.tensor([value for value, kind in zip(select(columns, 's').values(), kinds.values()) if kind == 'MONI'], dtype=dtype, device=device)\n\nAX = torch.tensor([value for value, kind in zip(select(columns, 'alphax').values(), kinds.values()) if kind == 'MONI'], dtype=dtype, device=device)\nBX = torch.tensor([value for value, kind in zip(select(columns, 'betax' ).values(), kinds.values()) if kind == 'MONI'], dtype=dtype, device=device)\nAY = torch.tensor([value for value, kind in zip(select(columns, 'alphay').values(), kinds.values()) if kind == 'MONI'], dtype=dtype, device=device)\nBY = torch.tensor([value for value, kind in zip(select(columns, 'betay' ).values(), kinds.values()) if kind == 'MONI'], dtype=dtype, device=device)\n\nMUX = torch.tensor([value for value, kind in zip(select(columns, 'psix').values(), kinds.values()) if kind == 'MONI'], dtype=dtype, device=device)\nMUY = torch.tensor([value for value, kind in zip(select(columns, 'psiy').values(), kinds.values()) if kind == 'MONI'], dtype=dtype, device=device)\n\nDQX = torch.tensor([value for value, kind in zip(select(columns, 'etax' ).values(), kinds.values()) if kind == 'MONI'], dtype=dtype, device=device)\nDPX = torch.tensor([value for value, kind in zip(select(columns, 'etaxp').values(), kinds.values()) if kind == 'MONI'], dtype=dtype, device=device)\nDQY = torch.tensor([value for value, kind in zip(select(columns, 'etay' ).values(), kinds.values()) if kind == 'MONI'], dtype=dtype, device=device)\nDPY = torch.tensor([value for value, kind in zip(select(columns, 'etayp').values(), kinds.values()) if kind == 'MONI'], dtype=dtype, device=device)\n```\n\n\n```python\n# Load and build ELEGANT lattice\n\nfrom model.command.external import load_lattice\n\npath = Path('ic.lte')\ndata = load_lattice(path)\n\nfrom model.library.line import Line\nfrom model.command.build import build\n\nring:Line = build('RING', 'ELEGANT', data)\n\n# Flatten lattice\n\nring.flatten()\n\n# Merge adjacent drifts (all drifts are renamed)\n\nring.merge()\n\n# Set linear flag in dipoles (dipoles are nonlinear by default)\n\nfor element in ring:\n    if element.kind == 'Dipole':\n        element.linear = True\n\n# Set number of elements of different kinds\n\nnb = ring.describe['BPM']\nnd = ring.describe['Dipole']\nnq = ring.describe['Quadrupole']\nns = ring.describe['Sextupole']\n\n# Describe\n\nprint(ring.describe)\n```\n\n    {'BPM': 16, 'Drift': 68, 'Quadrupole': 28, 'Dipole': 8, 'Sextupole': 16}\n\n\n\n```python\n# Plot 2D layout\n\nfrom model.command.layout import Layout\n\nline = ring.clone()\nline.ns = 0.01\nlayout = Layout(line)\n\n# Generate reference orbit\n\nx, y, _ = layout.orbit(flat=False, step=0.01, start=(0, 0))\n\n# Generate layout\n\nblocks, _ = layout.profile_2d(start=(0, 0), delta=1.0, linewidth=1.0, exclude=['Drift'])\n_, labels = layout.profile_2d(start=(0, 0), delta=1.0, linewidth=1.0, exclude=['Dipole', 'Quadrupole', 'Sextupole', 'Drift'], fontsize=8)\n\n# Plot\n\nfrom matplotlib import pyplot as plt\n\nplt.figure()\nplt.plot(x, y, color='black')\nfor block in blocks:\n    plt.errorbar(**block)\nfor label in labels:\n    plt.text(**label)\nplt.xlabel(r'x (m)')\nplt.ylabel(r'y (m)')\nplt.tight_layout()\nplt.show()\n```\n\n\n\n![png](demo/layout.png)\n\n\n\n\n```python\n# Tracking (single initial)\n\nx = torch.tensor([0.001, 0.0, 0.0, 0.0], dtype=dtype, device=device)\n\nprint(ring(x))\n```\n\n    tensor([-0.0045,  0.0091,  0.0000,  0.0000], dtype=torch.float64)\n\n\n\n```python\n# Tracking (vectorized map over initial values)\n\nxs = torch.tensor(2**10*[4*[0.0]], dtype=dtype, device=device)\n\nprint(torch.vmap(ring)(xs).shape)\n```\n\n    torch.Size([1024, 4])\n\n\n\n```python\n# Tracking (pass deviation parameters)\n\ndata = ring.data()\ndata['Q3D3']['kn'] = 0.25\n\nprint(ring(x, data=data))\n```\n\n    tensor([-0.0044,  0.0089,  0.0000,  0.0000], dtype=torch.float64)\n\n\n\n```python\n# Tracking (wrapper)\n\nfrom model.command.wrapper import group\n\nkn = torch.tensor([0.25], dtype=dtype, device=device)\nmapping, *_ = group(ring, 0, len(ring) - 1, ('kn', None, ['Q3D3'], None))\n\nprint(mapping(x, kn))\n```\n\n    tensor([-0.0044,  0.0089,  0.0000,  0.0000], dtype=torch.float64)\n\n\n\n```python\n# Differentiability with respect to initial state\n\nprint(torch.func.jacrev(mapping, 0)(x, kn))\nprint(torch.func.jacrev(mapping, 1)(x, kn))\n```\n\n    tensor([[-4.4739, -2.6532,  0.0000,  0.0000],\n            [ 8.9295,  5.0721,  0.0000,  0.0000],\n            [ 0.0000,  0.0000,  3.8653, -2.1508],\n            [ 0.0000,  0.0000,  5.5530, -2.8312]], dtype=torch.float64)\n    tensor([[ 0.0004],\n            [-0.0007],\n            [ 0.0000],\n            [ 0.0000]], dtype=torch.float64)\n\n\n\n```python\n# Differentiability with respect to deviation parameters\n\nprint(torch.func.jacrev(mapping, 1)(x, kn))\n```\n\n    tensor([[ 0.0004],\n            [-0.0007],\n            [ 0.0000],\n            [ 0.0000]], dtype=torch.float64)\n\n\n\n```python\n# Tunes\n\nfrom model.command.tune import tune\n\nnux, nuy = tune(ring, [], matched=True)\n\nprint(torch.allclose(NUX, nux, rtol=1.0E-12, atol=1.0E-12))\nprint(torch.allclose(NUY, nuy, rtol=1.0E-12, atol=1.0E-12))\n```\n\n    True\n    True\n\n\n\n```python\n# Parametric tunes\n\ndef fn(dp, kn, ks):\n    return tune(ring, [dp, kn, ks], ('dp', None, None, None), ('kn', ['Quadrupole'], None, None),  ('ks', ['Quadrupole'], None, None), matched=True)\n\n# Vectorizing map\n\ndp = torch.tensor([0.0], dtype=dtype, device=device)\n\nkn = torch.zeros(nq, dtype=dtype, device=device)\nks = torch.zeros(nq, dtype=dtype, device=device)\n\nprint(torch.vmap(fn, in_dims=(None, 0, None))(dp, torch.stack(2**10*[kn]), ks).shape)\n\n# Derivatives\n\nprint(torch.func.jacrev(fn, 0)(dp, kn, ks))\nprint(torch.func.jacrev(fn, 1)(dp, kn, ks))\n```\n\n    torch.Size([1024, 2])\n    tensor([[0.0284],\n            [0.7658]], dtype=torch.float64)\n    tensor([[ 0.0187,  0.0217,  0.0461,  0.0601,  0.0484,  0.0297,  0.0299,  0.0488,\n              0.0597,  0.0457,  0.0214,  0.0192,  0.0669,  0.0667,  0.0191,  0.0216,\n              0.0461,  0.0597,  0.0483,  0.0296,  0.0297,  0.0485,  0.0595,  0.0459,\n              0.0209,  0.0196,  0.0683,  0.0647],\n            [-0.0668, -0.0283, -0.0136, -0.0078, -0.0141, -0.0205, -0.0202, -0.0137,\n             -0.0079, -0.0137, -0.0282, -0.0661, -0.0250, -0.0251, -0.0663, -0.0282,\n             -0.0137, -0.0079, -0.0138, -0.0203, -0.0205, -0.0140, -0.0078, -0.0137,\n             -0.0288, -0.0659, -0.0250, -0.0258]], dtype=torch.float64)\n\n\n\n```python\n# Twiss parameters\n\nfrom model.command.twiss import twiss\n\nax, bx, ay, by = twiss(ring, [], alignment=False, matched=True, advance=True, full=False).T\n```\n\n\n```python\n# Compare Twiss parameters at BPMs\n\nax_bpm, bx_bpm, ay_bpm, by_bpm = torch.stack([torch.stack([*cs]) for (_, kind, *_), *cs in zip(ring.layout(), ax, bx, ay, by) if kind == 'BPM']).T\n\nprint(torch.allclose(AX, ax_bpm, rtol=1.0E-12, atol=1.0E-12))\nprint(torch.allclose(BX, bx_bpm, rtol=1.0E-12, atol=1.0E-12))\nprint(torch.allclose(AY, ay_bpm, rtol=1.0E-12, atol=1.0E-12))\nprint(torch.allclose(BY, by_bpm, rtol=1.0E-12, atol=1.0E-12))\n```\n\n    True\n    True\n    True\n    True\n\n\n\n```python\n# Plot beta functions\n\nfrom matplotlib.patches import Rectangle\n\nlayout = Layout(ring)\n_, _, lengths, *_ = layout.slicing_table()\nrectangles = layout.profile_1d(scale=0.75, shift=0.0, text=False,  exclude=['BPM', 'Drift'])\n\nplt.figure(figsize=(16, 4))\nplt.errorbar(S.cpu().numpy(), BX.cpu().numpy(), fmt=' ', color='red', alpha=0.75, marker='x')\nplt.errorbar(ring.locations().cpu().numpy(), bx.cpu().numpy(), fmt='-', color='red', alpha=0.75)\nplt.errorbar(S.cpu().numpy(), BY.cpu().numpy(), fmt=' ', color='blue', alpha=0.75, marker='x')\nplt.errorbar(ring.locations().cpu().numpy(), by.cpu().numpy(), fmt='-', color='blue', alpha=0.75)\nfor rectangle in rectangles:\n    plt.gca().add_patch(Rectangle(**rectangle))\nplt.ylim(-0.5, 5.0)\nplt.tight_layout()\nplt.show()\n```\n\n\n\n![png](demo/beta.png)\n\n\n\n\n```python\n# Parametric twiss (lattice start location)\n\ndef fn(dp, kn, ks):\n    return twiss(ring, [dp, kn, ks], ('dp', None, None, None), ('kn', ['Quadrupole'], None, None),  ('ks', ['Quadrupole'], None, None), advance=False, matched=True)\n\n# Vectorizing map\n\ndp = torch.tensor([0.0], dtype=dtype, device=device)\n\nkn = torch.zeros(nq, dtype=dtype, device=device)\nks = torch.zeros(nq, dtype=dtype, device=device)\n\nprint(torch.vmap(fn, in_dims=(None, 0, None))(dp, torch.stack(2**10*[kn]), ks).shape)\n\n# Derivatives\n\nprint(torch.func.jacrev(fn, 0)(dp, kn, ks))\nprint(torch.func.jacrev(fn, 1)(dp, kn, ks))\n```\n\n    torch.Size([1024, 4])\n    tensor([[-222.0538],\n            [-110.7909],\n            [  27.9708],\n            [ -13.9537]], dtype=torch.float64)\n    tensor([[-0.0893,  0.4014,  1.2554, -0.9068, -1.5491, -0.9866, -0.6147, -0.7071,\n             -1.9186,  0.2045, -0.1659,  0.4221,  1.9239,  2.2147,  0.5854, -0.4487,\n             -0.4684, -1.9595, -0.0112, -0.2204, -0.9209, -1.5967, -0.0541,  1.5081,\n              0.5988, -0.3222, -0.4638,  0.8415],\n            [ 0.0182,  0.1496,  0.5733, -0.6730, -0.7756, -0.5178, -0.4104, -0.5325,\n             -0.9651,  0.2699, -0.0145,  0.1734,  0.9035,  1.1954,  0.3374, -0.2984,\n             -0.4055, -1.0062,  0.1660, -0.0147, -0.4518, -0.8235,  0.1829,  0.8281,\n              0.3562, -0.2319, -0.4848,  0.2358],\n            [ 1.6183, -0.0219, -0.2993, -0.0049,  0.3526,  0.2005, -0.5750, -0.4693,\n              0.1832,  0.0065, -0.5948, -2.2667, -0.8472,  0.8119,  2.2819,  0.7206,\n              0.0759, -0.2127,  0.4684,  0.6358, -0.0792, -0.2893, -0.0394,  0.3584,\n              0.2001, -1.2685, -0.6598,  0.2486],\n            [-0.7202,  0.1496,  0.2258, -0.0342, -0.1608, -0.0228,  0.3921,  0.2789,\n             -0.1348,  0.0622,  0.4565,  1.3632,  0.4758, -0.5180, -1.3618, -0.3315,\n              0.0206,  0.1024, -0.2828, -0.3332,  0.1447,  0.2237, -0.0135, -0.1681,\n              0.0184,  1.0088,  0.4647, -0.2653]], dtype=torch.float64)\n\n\n\n```python\n# Phase advance\n\nfrom model.command.advance import advance\n\nmux, muy =  advance(ring, [], alignment=False, matched=True).T\n```\n\n\n```python\n# Total tunes\n\nprint(torch.allclose(mux.sum()/(2*torch.pi), torch.tensor(parameters['nux'], dtype=torch.float64), rtol=1.0E-12, atol=1.0E-12))\nprint(torch.allclose(muy.sum()/(2*torch.pi), torch.tensor(parameters['nuy'], dtype=torch.float64), rtol=1.0E-12, atol=1.0E-12))\n```\n\n    True\n    True\n\n\n\n```python\n# Plot accumulated phase advance\n\nmux_sum = torch.cat([torch.tensor([0.0], dtype=torch.float64), mux.cumsum(-1)])\nmuy_sum = torch.cat([torch.tensor([0.0], dtype=torch.float64), muy.cumsum(-1)])\n\nlayout = Layout(ring)\n_, _, lengths, *_ = layout.slicing_table()\nrectangles = layout.profile_1d(scale=5.0, shift=-2.5, text=False,  exclude=['BPM', 'Drift'])\n\nplt.figure(figsize=(16, 4))\nplt.errorbar(S.cpu().numpy(), MUX.cpu().numpy(), fmt=' ', color='red', alpha=0.75, marker='x')\nplt.errorbar(torch.cat([ring.locations(), ring.length.unsqueeze(-1)]).cpu().numpy(), mux_sum.cpu().numpy(), fmt='-', color='red', alpha=0.75)\nplt.errorbar(S.cpu().numpy(), MUY.cpu().numpy(), fmt=' ', color='blue', alpha=0.75, marker='x')\nplt.errorbar(torch.cat([ring.locations(), ring.length.unsqueeze(-1)]).cpu().numpy(), muy_sum.cpu().numpy(), fmt='-', color='blue', alpha=0.75)\nfor rectangle in rectangles:\n    plt.gca().add_patch(Rectangle(**rectangle))\nplt.ylim(-5.0, 35.0)\nplt.tight_layout()\nplt.show()\n```\n\n\n\n![png](demo/advance.png)\n\n\n\n\n```python\n# Parametric advance\n\ndef fn(dp, kn, ks):\n    return advance(ring, [dp, kn, ks], ('dp', None, None, None), ('kn', ['Quadrupole'], None, None),  ('ks', ['Quadrupole'], None, None), matched=True)\n\n# Vectorizing map\n\ndp = torch.tensor([0.0], dtype=dtype, device=device)\n\nkn = torch.zeros(nq, dtype=dtype, device=device)\nks = torch.zeros(nq, dtype=dtype, device=device)\n\nprint(torch.vmap(fn, in_dims=(None, 0, None))(dp, torch.stack(2**10*[kn]), ks).shape)\n\n# Derivatives\n\nprint(torch.func.jacrev(fn, 0)(dp, kn, ks).shape)\nprint(torch.func.jacrev(fn, 1)(dp, kn, ks).shape)\n```\n\n    torch.Size([1024, 136, 2])\n    torch.Size([136, 2, 1])\n    torch.Size([136, 2, 28])\n\n\n\n```python\n# Dispersion\n\nfrom model.command.orbit import dispersion\n\nguess = torch.tensor(4*[0.0], dtype=dtype, device=device)\ndqx, dpx, dqy, dpy = dispersion(ring, guess, [], advance=True, full=False)\n```\n\n\n```python\n# Compare dispersion at BPMs\n\ndqx_bpm, dpx_bpm, dqy_bpm, dpy_bpm = torch.stack([torch.stack([*cs]) for (_, kind, *_), *cs in zip(ring.layout(), dqx, dpx, dqy, dpy) if kind == 'BPM']).T\n\nprint(torch.allclose(DQX, dqx_bpm, rtol=1.0E-12, atol=1.0E-12))\nprint(torch.allclose(DPX, dpx_bpm, rtol=1.0E-12, atol=1.0E-12))\nprint(torch.allclose(DQY, dqy_bpm, rtol=1.0E-12, atol=1.0E-12))\nprint(torch.allclose(DPY, dpy_bpm, rtol=1.0E-12, atol=1.0E-12))\n```\n\n    True\n    True\n    True\n    True\n\n\n\n```python\n# Plot dispersion\n\nlayout = Layout(ring)\n_, _, lengths, *_ = layout.slicing_table()\nrectangles = layout.profile_1d(scale=0.25, shift=-0.25, text=False,  exclude=['BPM', 'Drift'])\n\nplt.figure(figsize=(16, 4))\nplt.errorbar(S.cpu().numpy(), DQX.cpu().numpy(), fmt=' ', color='red', alpha=0.75, marker='x')\nplt.errorbar(ring.locations().cpu().numpy(), dqx.cpu().numpy(), fmt='-', color='red', alpha=0.75)\nfor rectangle in rectangles:\n    plt.gca().add_patch(Rectangle(**rectangle))\nplt.ylim(-0.5, 1.0)\nplt.tight_layout()\nplt.show()\n```\n\n\n\n![png](demo/dispersion.png)\n\n\n\n\n```python\n# Parametric dispersion\n\ndef fn(kn, ks):\n    return dispersion(ring, guess, [kn, ks], ('kn', ['Quadrupole'], None, None),  ('ks', ['Quadrupole'], None, None), advance=False)\n\n# Vectorizing map\n\nkn = torch.zeros(nq, dtype=dtype, device=device)\nks = torch.zeros(nq, dtype=dtype, device=device)\n\nprint(torch.vmap(fn, in_dims=(0, None))(torch.stack(2**10*[kn]), ks).shape)\n\n# Derivatives\n\nprint(torch.func.jacrev(fn, 0)(kn, ks))\nprint(torch.func.jacrev(fn, 1)(kn, ks))\n```\n\n    torch.Size([1024, 4])\n    tensor([[-1.1162e-03, -8.2786e-04, -2.4801e-03,  1.9629e-01,  2.7563e-05,\n              9.5180e-06, -2.3718e-05, -4.7673e-05, -1.0284e-01,  2.2051e-03,\n              6.9025e-04,  1.3755e-03,  5.3769e-03,  5.7391e-03,  1.5789e-03,\n              3.3797e-04,  1.3670e-03,  8.2397e-02, -6.1463e-05, -3.5245e-05,\n             -1.4306e-05, -1.4063e-05,  3.4270e-01, -2.6831e-03, -9.3952e-04,\n             -6.7697e-04, -3.2418e-03, -4.2953e-03],\n            [ 1.6338e-03,  1.6955e-03,  4.8680e-03, -5.0176e-01, -9.1579e-06,\n              8.7653e-06,  6.6251e-05,  1.2097e-04,  2.7302e-02, -3.4149e-03,\n             -9.8856e-04, -2.7727e-03, -1.0457e-02, -1.0501e-02, -2.7973e-03,\n             -9.4603e-04, -3.3137e-03,  1.1249e-02,  9.2654e-05,  5.0532e-05,\n              8.0420e-06, -3.5286e-06, -5.1448e-01,  4.8677e-03,  1.6412e-03,\n              1.7617e-03,  7.6393e-03,  9.0321e-03],\n            [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00],\n            [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00]], dtype=torch.float64)\n    tensor([[ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00],\n            [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n              0.0000e+00,  0.0000e+00,  0.0000e+00],\n            [ 6.0868e-03, -1.9121e-03, -1.6954e-03,  2.1783e-01,  1.3196e-04,\n              1.1426e-04,  4.9850e-05,  1.7995e-05, -3.0953e-01, -2.5745e-03,\n             -1.3526e-03,  1.0628e-03,  1.8173e-04, -7.3803e-03, -6.3713e-03,\n              2.3662e-03,  2.6927e-03, -7.3148e-02, -9.2005e-05, -9.1568e-05,\n             -6.9841e-05, -5.2553e-05,  1.9908e-01,  3.1329e-03,  2.0529e-03,\n             -3.7047e-03, -3.4655e-03,  5.5212e-03],\n            [ 3.3306e-03, -9.0348e-04, -6.1116e-04,  1.6731e-01,  7.2547e-05,\n              5.7852e-05,  1.3325e-05, -6.7422e-06, -2.0683e-01, -1.2303e-03,\n             -5.0343e-04, -2.6321e-04, -9.7110e-04, -4.6478e-03, -3.8785e-03,\n              1.3020e-03,  1.3295e-03, -9.1948e-02, -5.6294e-05, -5.2568e-05,\n             -3.2302e-05, -1.9640e-05,  1.5820e-01,  1.7351e-03,  1.0243e-03,\n             -1.4231e-03, -1.0576e-03,  3.9832e-03]], dtype=torch.float64)\n\n\n\n```python\n# Modify lattice for ORM computation\n\nfrom model.library.corrector import Corrector\n\n# Clone lattice\n\ncopy = ring.clone()\n\n# Split BPMs and change lattice start\n\ncopy.split((None, ['BPM'], None, None))\ncopy.roll(1)\n\n# Insert correctors into quadrupoles and dipoles\n\nncq = 2**2\nncd = 2**4\n\nfor name in [name for name, kind, *_ in copy.layout() if kind == 'Quadrupole']:\n    corrector = Corrector(f'{name}_CXY', factor=1/(ncq - 1))\n    copy.split((ncq, None, [name], None), paste=[corrector])\n\nfor name in [name for name, kind, *_ in copy.layout() if kind == 'Dipole']:\n    corrector = Corrector(f'{name}_CXY', factor=1/(ncd - 1))\n    copy.split((ncd, None, [name], None), paste=[corrector])\n\n# Splice lattice\n\ncopy.splice()\n\n# Print sequence between first two BPMs\n\nline, *_ = copy\nprint(line)\n```\n\n    BPM(name=\"BPM05\", direction=\"inverse\", dp=0.0)\n    Drift(name=\"DR001\", length=0.116715, dp=0.0, exact=False, ns=1, order=0)\n    Quadrupole(name=\"Q3D3\", length=0.0491425, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n    Corrector(name=\"Q3D3_CXY\", cx=0.0, cy=0.0, factor=0.3333333333333333, dp=0.0)\n    Quadrupole(name=\"Q3D3\", length=0.0491425, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n    Corrector(name=\"Q3D3_CXY\", cx=0.0, cy=0.0, factor=0.3333333333333333, dp=0.0)\n    Quadrupole(name=\"Q3D3\", length=0.0491425, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n    Corrector(name=\"Q3D3_CXY\", cx=0.0, cy=0.0, factor=0.3333333333333333, dp=0.0)\n    Quadrupole(name=\"Q3D3\", length=0.0491425, kn=-8.426928737999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n    Drift(name=\"DR002\", length=0.9684299999999999, dp=0.0, exact=False, ns=1, order=0)\n    Quadrupole(name=\"Q3D2\", length=0.0491425, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n    Corrector(name=\"Q3D2_CXY\", cx=0.0, cy=0.0, factor=0.3333333333333333, dp=0.0)\n    Quadrupole(name=\"Q3D2\", length=0.0491425, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n    Corrector(name=\"Q3D2_CXY\", cx=0.0, cy=0.0, factor=0.3333333333333333, dp=0.0)\n    Quadrupole(name=\"Q3D2\", length=0.0491425, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n    Corrector(name=\"Q3D2_CXY\", cx=0.0, cy=0.0, factor=0.3333333333333333, dp=0.0)\n    Quadrupole(name=\"Q3D2\", length=0.0491425, kn=-2.695188250999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n    Drift(name=\"DR003\", length=0.221715, dp=0.0, exact=False, ns=1, order=0)\n    Quadrupole(name=\"Q3F2\", length=0.045, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n    Corrector(name=\"Q3F2_CXY\", cx=0.0, cy=0.0, factor=0.3333333333333333, dp=0.0)\n    Quadrupole(name=\"Q3F2\", length=0.045, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n    Corrector(name=\"Q3F2_CXY\", cx=0.0, cy=0.0, factor=0.3333333333333333, dp=0.0)\n    Quadrupole(name=\"Q3F2\", length=0.045, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n    Corrector(name=\"Q3F2_CXY\", cx=0.0, cy=0.0, factor=0.3333333333333333, dp=0.0)\n    Quadrupole(name=\"Q3F2\", length=0.045, kn=13.544085930000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n    Drift(name=\"DR004\", length=0.065, dp=0.0, exact=False, ns=1, order=0)\n    BPM(name=\"BPM07\", direction=\"forward\", dp=0.0)\n\n\n\n```python\n# Compute and plot ORM\n\nfrom model.command.orbit import ORM\n\norm = ORM(copy, guess, [])\nprint(orm.shape)\n\ndata = orm.clone()\ndata[data==0.0] = torch.nan\n\nplt.figure(figsize=(10, 10))\nimg = plt.imshow(data.cpu().numpy(), cmap='magma', interpolation='nearest')\ncax = plt.gcf().add_axes([plt.gca().get_position().x1 + 0.01, plt.gca().get_position().y0, 0.02, plt.gca().get_position().height])\nplt.colorbar(img, cax=cax)\nplt.show()\n```\n\n    torch.Size([32, 72])\n\n\n\n\n![png](demo/orm.png)\n\n\n\n\n```python\n# Parametric ORM\n\ndef fn(dp, kn, ks):\n    return ORM(copy, guess, [dp, kn, ks], ('dp', None, None, None), ('kn', ['Quadrupole'], None, None),  ('ks', ['Quadrupole'], None, None))\n\ndp = torch.tensor([0.0], dtype=dtype, device=device)\nkn = torch.zeros(nq, dtype=dtype, device=device)\nks = torch.zeros(nq, dtype=dtype, device=device)\n\nprint(torch.func.jacrev(fn, 0)(dp, kn, ks).shape)\nprint(torch.func.jacrev(fn, 1)(dp, kn, ks).shape)\n```\n\n    torch.Size([32, 72, 1])\n    torch.Size([32, 72, 28])\n\n\n# Related projects\n\n[https://i-a-morozov.github.io/twiss/](https://i-a-morozov.github.io/twiss/)\n\n[https://i-a-morozov.github.io/ndmap/](https://i-a-morozov.github.io/ndmap/)\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fi-a-morozov%2Fmodel","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fi-a-morozov%2Fmodel","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fi-a-morozov%2Fmodel/lists"}