mbsolver 0.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
mbsolver-0.2.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Vahid Anari
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,127 @@
1
+ Metadata-Version: 2.4
2
+ Name: mbsolver
3
+ Version: 0.2.0
4
+ Summary: Fast Maxwell-Bloch solvers for masers, superradiance and coherent pulse propagation
5
+ Author-email: Vahid Anari <vahid.anari8@gmail.com>
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/vahid-anari/mbsolver
8
+ Project-URL: Repository, https://github.com/vahid-anari/mbsolver
9
+ Project-URL: Issues, https://github.com/vahid-anari/mbsolver/issues
10
+ Keywords: maxwell-bloch,maser,superradiance,physics,astrophysics,numba,simulation
11
+ Classifier: Development Status :: 3 - Alpha
12
+ Classifier: Intended Audience :: Science/Research
13
+ Classifier: Operating System :: OS Independent
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Programming Language :: Python :: 3.13
19
+ Classifier: Topic :: Scientific/Engineering :: Physics
20
+ Classifier: Topic :: Scientific/Engineering :: Astronomy
21
+ Requires-Python: >=3.10
22
+ Description-Content-Type: text/markdown
23
+ License-File: LICENSE
24
+ Requires-Dist: numpy>=1.24
25
+ Requires-Dist: numba>=0.58
26
+ Provides-Extra: dev
27
+ Requires-Dist: build; extra == "dev"
28
+ Requires-Dist: twine; extra == "dev"
29
+ Requires-Dist: pytest; extra == "dev"
30
+ Dynamic: license-file
31
+
32
+ # mbsolver
33
+
34
+ Fast Maxwell-Bloch solvers in Python (numpy + numba) for astrophysical masers,
35
+ superradiance and coherent pulse propagation. Use them from scripts,
36
+ notebooks or cluster jobs; the
37
+ [Maxwell Bloch Studio](https://github.com/vahid-anari/maxwell-bloch-studio)
38
+ desktop app runs on the same code.
39
+
40
+ ```
41
+ pip install mbsolver
42
+ ```
43
+
44
+ ## Two-level model in one dimension
45
+
46
+ Dimensionless equations for the population inversion `w`, the coherence `R`
47
+ and the field `A` along a medium `0 ≤ z ≤ 1`:
48
+
49
+ ```
50
+ dw/dt = A R − (w − w0) / T1 + Λ(t)
51
+ dR/dt = −A w − (R − R0) / T2
52
+ dA/dz = −η R
53
+ ```
54
+
55
+ The field `A(t)` entering at `z = 0` is given; population and coherence are
56
+ advanced with fourth-order Runge-Kutta in time, the field step by step in `z`.
57
+
58
+ ```python
59
+ import matplotlib.pyplot as plt
60
+ from mbsolver import Pulse, Steps
61
+ from mbsolver.two_level import solve_1d
62
+
63
+ res = solve_1d(
64
+ t_max=300, nt=3001, nz=121, eta=40.0,
65
+ theta0=0.003, # or w0=..., R0=...
66
+ pump=[Pulse(a=[2e-3, 1e-3], x0=[60, 180], w=[15, 25])],
67
+ field_in=Pulse(a=1e-3, x0=30, w=5), # A(t) at z = 0
68
+ t1=Steps(y=[80, 40, 60], x0=[100, 200], w=[10, 0]),
69
+ t2=5.0,
70
+ n_planes=7, # stored z planes, 0 ... 1
71
+ )
72
+
73
+ plt.plot(res.time, res.intensity[-1]) # intensity leaving the medium
74
+ plt.show()
75
+ ```
76
+
77
+ ### Arguments
78
+
79
+ | Argument | Meaning |
80
+ |----------|---------|
81
+ | `t_max`, `nt` | time grid from 0 to `t_max` with `nt` points |
82
+ | `nz` | number of z points |
83
+ | `eta` | coupling coefficient η |
84
+ | `theta0`, or `w0` and `R0` | initial state (`w0 = cos θ0`, `R0 = sin θ0`); the medium also relaxes towards it |
85
+ | `t1`, `t2` | relaxation times: a number or `Steps` (> 0, `math.inf` for none) |
86
+ | `pump` | Λ(t): a `Pulse`, a number, or a list of them (summed) |
87
+ | `field_in` | A(t) at z = 0: a `Pulse`, a number, or `None` |
88
+ | `n_planes` | number of evenly spaced z planes returned (default 2: entrance and exit) |
89
+
90
+ ### Result
91
+
92
+ `res.time`, `res.z` (positions of the stored planes), `res.intensity` and
93
+ `res.w` (shape `(n_planes, nt)`), `res.pump`, `res.field_in`, `res.t1`, `res.t2`.
94
+
95
+ ## Profiles
96
+
97
+ `Pulse(a, x0, w)` is a sum of sech² pulses, `a_i sech²((t − x0_i) / w_i)`.
98
+ Give `wl` and `wr` instead of `w` for a different width on each side.
99
+
100
+ `Steps(y, x0, w)` goes from level `y[0]` to `y[1]` around `x0[0]`, then to
101
+ `y[2]` around `x0[1]`, and so on, each step smoothed as `(1 + tanh)/2` with
102
+ width `w` (0 for a sharp step). `x0` and `w` have one value fewer than `y`.
103
+
104
+ A zero width gives a single point (pulse) or a sharp step; an infinite width
105
+ a constant. `profile.evaluate(t)` returns the values on any grid.
106
+
107
+ ## Tests
108
+
109
+ ```
110
+ pip install -e ".[dev]"
111
+ pytest
112
+ ```
113
+
114
+ The tests include a reference run that must match the solver of the Maxwell
115
+ Bloch app (version 8.13) exactly, and physics checks (Bloch rotation at the
116
+ entrance, relaxation balance, uncoupled propagation).
117
+
118
+ ## Earlier name
119
+
120
+ Up to 0.1.0 this package was called `maxbloch`; 0.2.0 is the same code under
121
+ the new name. `maxbloch` stays on PyPI unchanged (0.1.1 only adds a
122
+ deprecation warning) but gets no updates; change `maxbloch` to `mbsolver` in
123
+ your imports.
124
+
125
+ ## License
126
+
127
+ MIT
@@ -0,0 +1,96 @@
1
+ # mbsolver
2
+
3
+ Fast Maxwell-Bloch solvers in Python (numpy + numba) for astrophysical masers,
4
+ superradiance and coherent pulse propagation. Use them from scripts,
5
+ notebooks or cluster jobs; the
6
+ [Maxwell Bloch Studio](https://github.com/vahid-anari/maxwell-bloch-studio)
7
+ desktop app runs on the same code.
8
+
9
+ ```
10
+ pip install mbsolver
11
+ ```
12
+
13
+ ## Two-level model in one dimension
14
+
15
+ Dimensionless equations for the population inversion `w`, the coherence `R`
16
+ and the field `A` along a medium `0 ≤ z ≤ 1`:
17
+
18
+ ```
19
+ dw/dt = A R − (w − w0) / T1 + Λ(t)
20
+ dR/dt = −A w − (R − R0) / T2
21
+ dA/dz = −η R
22
+ ```
23
+
24
+ The field `A(t)` entering at `z = 0` is given; population and coherence are
25
+ advanced with fourth-order Runge-Kutta in time, the field step by step in `z`.
26
+
27
+ ```python
28
+ import matplotlib.pyplot as plt
29
+ from mbsolver import Pulse, Steps
30
+ from mbsolver.two_level import solve_1d
31
+
32
+ res = solve_1d(
33
+ t_max=300, nt=3001, nz=121, eta=40.0,
34
+ theta0=0.003, # or w0=..., R0=...
35
+ pump=[Pulse(a=[2e-3, 1e-3], x0=[60, 180], w=[15, 25])],
36
+ field_in=Pulse(a=1e-3, x0=30, w=5), # A(t) at z = 0
37
+ t1=Steps(y=[80, 40, 60], x0=[100, 200], w=[10, 0]),
38
+ t2=5.0,
39
+ n_planes=7, # stored z planes, 0 ... 1
40
+ )
41
+
42
+ plt.plot(res.time, res.intensity[-1]) # intensity leaving the medium
43
+ plt.show()
44
+ ```
45
+
46
+ ### Arguments
47
+
48
+ | Argument | Meaning |
49
+ |----------|---------|
50
+ | `t_max`, `nt` | time grid from 0 to `t_max` with `nt` points |
51
+ | `nz` | number of z points |
52
+ | `eta` | coupling coefficient η |
53
+ | `theta0`, or `w0` and `R0` | initial state (`w0 = cos θ0`, `R0 = sin θ0`); the medium also relaxes towards it |
54
+ | `t1`, `t2` | relaxation times: a number or `Steps` (> 0, `math.inf` for none) |
55
+ | `pump` | Λ(t): a `Pulse`, a number, or a list of them (summed) |
56
+ | `field_in` | A(t) at z = 0: a `Pulse`, a number, or `None` |
57
+ | `n_planes` | number of evenly spaced z planes returned (default 2: entrance and exit) |
58
+
59
+ ### Result
60
+
61
+ `res.time`, `res.z` (positions of the stored planes), `res.intensity` and
62
+ `res.w` (shape `(n_planes, nt)`), `res.pump`, `res.field_in`, `res.t1`, `res.t2`.
63
+
64
+ ## Profiles
65
+
66
+ `Pulse(a, x0, w)` is a sum of sech² pulses, `a_i sech²((t − x0_i) / w_i)`.
67
+ Give `wl` and `wr` instead of `w` for a different width on each side.
68
+
69
+ `Steps(y, x0, w)` goes from level `y[0]` to `y[1]` around `x0[0]`, then to
70
+ `y[2]` around `x0[1]`, and so on, each step smoothed as `(1 + tanh)/2` with
71
+ width `w` (0 for a sharp step). `x0` and `w` have one value fewer than `y`.
72
+
73
+ A zero width gives a single point (pulse) or a sharp step; an infinite width
74
+ a constant. `profile.evaluate(t)` returns the values on any grid.
75
+
76
+ ## Tests
77
+
78
+ ```
79
+ pip install -e ".[dev]"
80
+ pytest
81
+ ```
82
+
83
+ The tests include a reference run that must match the solver of the Maxwell
84
+ Bloch app (version 8.13) exactly, and physics checks (Bloch rotation at the
85
+ entrance, relaxation balance, uncoupled propagation).
86
+
87
+ ## Earlier name
88
+
89
+ Up to 0.1.0 this package was called `maxbloch`; 0.2.0 is the same code under
90
+ the new name. `maxbloch` stays on PyPI unchanged (0.1.1 only adds a
91
+ deprecation warning) but gets no updates; change `maxbloch` to `mbsolver` in
92
+ your imports.
93
+
94
+ ## License
95
+
96
+ MIT
@@ -0,0 +1,47 @@
1
+ [build-system]
2
+ requires = ["setuptools>=77"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "mbsolver"
7
+ version = "0.2.0"
8
+ description = "Fast Maxwell-Bloch solvers for masers, superradiance and coherent pulse propagation"
9
+ readme = "README.md"
10
+ requires-python = ">=3.10"
11
+ license = "MIT"
12
+ license-files = ["LICENSE"]
13
+ authors = [{ name = "Vahid Anari", email = "vahid.anari8@gmail.com" }]
14
+ keywords = [
15
+ "maxwell-bloch", "maser", "superradiance", "physics",
16
+ "astrophysics", "numba", "simulation",
17
+ ]
18
+ classifiers = [
19
+ "Development Status :: 3 - Alpha",
20
+ "Intended Audience :: Science/Research",
21
+ "Operating System :: OS Independent",
22
+ "Programming Language :: Python :: 3",
23
+ "Programming Language :: Python :: 3.10",
24
+ "Programming Language :: Python :: 3.11",
25
+ "Programming Language :: Python :: 3.12",
26
+ "Programming Language :: Python :: 3.13",
27
+ "Topic :: Scientific/Engineering :: Physics",
28
+ "Topic :: Scientific/Engineering :: Astronomy",
29
+ ]
30
+ dependencies = [
31
+ "numpy>=1.24",
32
+ "numba>=0.58",
33
+ ]
34
+
35
+ [project.optional-dependencies]
36
+ dev = ["build", "twine", "pytest"]
37
+
38
+ [project.urls]
39
+ Homepage = "https://github.com/vahid-anari/mbsolver"
40
+ Repository = "https://github.com/vahid-anari/mbsolver"
41
+ Issues = "https://github.com/vahid-anari/mbsolver/issues"
42
+
43
+ [tool.setuptools.packages.find]
44
+ where = ["src"]
45
+
46
+ [tool.pytest.ini_options]
47
+ testpaths = ["tests"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,17 @@
1
+ """Maxwell-Bloch solvers.
2
+
3
+ Profiles (pump pulses, boundary fields, relaxation-time schedules) are built
4
+ with :class:`Pulse` and :class:`Steps`; the solvers live in model subpackages,
5
+ e.g. ``mbsolver.two_level.solve_1d``.
6
+ """
7
+
8
+ from importlib.metadata import PackageNotFoundError, version
9
+
10
+ from mbsolver.profiles import Pulse, Steps
11
+
12
+ try:
13
+ __version__ = version("mbsolver")
14
+ except PackageNotFoundError: # running from a source tree without install
15
+ __version__ = "0+unknown"
16
+
17
+ __all__ = ["Pulse", "Steps", "__version__"]
@@ -0,0 +1,16 @@
1
+ """Simulation grids."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import numpy as np
6
+ from numba import njit
7
+
8
+
9
+ @njit(inline="always", nogil=True, cache=True)
10
+ def grids_1d(t_max: float, nt: int, nz: int):
11
+ """Return ``(t, dt, z, dz)``: ``nt`` times from 0 to ``t_max``, ``nz`` positions from 0 to 1."""
12
+ t = np.linspace(0.0, t_max, nt)
13
+ dt = t_max / (nt - 1)
14
+ z = np.linspace(0.0, 1.0, nz)
15
+ dz = 1.0 / (nz - 1)
16
+ return t, dt, z, dz
@@ -0,0 +1,237 @@
1
+ """Time profiles: sums of sech² pulses and smoothed steps.
2
+
3
+ ``Pulse`` describes pump pulses and boundary fields, ``Steps`` describes
4
+ quantities that change between levels, such as relaxation times. Both are
5
+ evaluated on a grid with numba.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ from dataclasses import dataclass
11
+ from typing import Optional, Sequence, Union
12
+
13
+ import numpy as np
14
+ from numba import njit
15
+
16
+ ArrayLike = Union[float, Sequence[float], np.ndarray]
17
+
18
+
19
+ def _array(name: str, values: Optional[ArrayLike]) -> Optional[np.ndarray]:
20
+ """Return ``values`` as a 1-D float array (``None`` stays ``None``)."""
21
+ if values is None:
22
+ return None
23
+ arr = np.atleast_1d(np.asarray(values, dtype=np.float64))
24
+ if arr.ndim != 1:
25
+ raise ValueError(f"{name} must be a number or a 1-D sequence")
26
+ return arr
27
+
28
+
29
+ def _widths(n: int, w, wl, wr) -> tuple[bool, np.ndarray, np.ndarray, np.ndarray]:
30
+ """Check the width arguments and return ``(symmetric, w, wl, wr)`` arrays of length ``n``."""
31
+ w, wl, wr = _array("w", w), _array("wl", wl), _array("wr", wr)
32
+ if w is not None:
33
+ if wl is not None or wr is not None:
34
+ raise ValueError("give either w, or wl and wr, not both")
35
+ symmetric, wl, wr = True, w, w
36
+ elif wl is not None and wr is not None:
37
+ symmetric, w = False, wl
38
+ else:
39
+ raise ValueError("give either w, or both wl and wr")
40
+ for name, arr in (("w", w), ("wl", wl), ("wr", wr)):
41
+ if arr.shape[0] != n:
42
+ raise ValueError(f"{name} must have {n} values, got {arr.shape[0]}")
43
+ if np.any(arr < 0):
44
+ raise ValueError(f"{name} must not be negative")
45
+ return symmetric, w, wl, wr
46
+
47
+
48
+ @dataclass(frozen=True, init=False)
49
+ class Pulse:
50
+ """A sum of sech² pulses.
51
+
52
+ ``y(x) = sum_i a_i * sech²((x - x0_i) / w_i)``; with ``wl`` and ``wr``
53
+ instead of ``w``, each pulse has its own width on each side of ``x0_i``.
54
+ A zero width gives ``a_i`` only at ``x0_i``; an infinite width a constant.
55
+
56
+ Args:
57
+ a: Amplitudes.
58
+ x0: Centres.
59
+ w: Widths (symmetric pulses).
60
+ wl, wr: Left and right widths (asymmetric pulses).
61
+
62
+ Example:
63
+ ``Pulse(a=[0.1, 0.05], x0=[50, 120], w=[10, 20])``
64
+ """
65
+
66
+ a: np.ndarray
67
+ x0: np.ndarray
68
+ w: np.ndarray
69
+ wl: np.ndarray
70
+ wr: np.ndarray
71
+ symmetric: bool
72
+
73
+ def __init__(self, a: ArrayLike, x0: ArrayLike, w: Optional[ArrayLike] = None,
74
+ wl: Optional[ArrayLike] = None, wr: Optional[ArrayLike] = None) -> None:
75
+ a, x0 = _array("a", a), _array("x0", x0)
76
+ n = a.shape[0]
77
+ if x0.shape[0] != n:
78
+ raise ValueError(f"x0 must have {n} values (one per amplitude), got {x0.shape[0]}")
79
+ symmetric, w, wl, wr = _widths(n, w, wl, wr)
80
+ for name, value in (("a", a), ("x0", x0), ("w", w), ("wl", wl), ("wr", wr),
81
+ ("symmetric", symmetric)):
82
+ object.__setattr__(self, name, value)
83
+
84
+ def evaluate(self, x: ArrayLike) -> np.ndarray:
85
+ """Return the profile on the coordinates ``x``."""
86
+ x = np.asarray(x, dtype=np.float64)
87
+ return pulse_func(self.symmetric, x, self.a, self.x0, self.w, self.wl, self.wr)
88
+
89
+
90
+ @dataclass(frozen=True, init=False)
91
+ class Steps:
92
+ """Smoothed steps between levels.
93
+
94
+ Starts at ``y[0]`` and changes to ``y[i]`` around ``x0[i-1]``:
95
+ ``y(x) = y_0 + sum_i (y_i - y_{i-1}) * (1 + tanh((x - x0_i) / w_i)) / 2``.
96
+ A zero width gives a sharp step; one level gives a constant.
97
+
98
+ Args:
99
+ y: Levels (``N`` values).
100
+ x0: Step positions (``N - 1`` values).
101
+ w: Step widths (``N - 1`` values, symmetric steps).
102
+ wl, wr: Left and right widths (asymmetric steps).
103
+
104
+ Example:
105
+ ``Steps(y=[0.5, 0.2], x0=[100], w=[10])``: 0.5 before t = 100, 0.2 after.
106
+ """
107
+
108
+ y: np.ndarray
109
+ x0: np.ndarray
110
+ w: np.ndarray
111
+ wl: np.ndarray
112
+ wr: np.ndarray
113
+ symmetric: bool
114
+
115
+ def __init__(self, y: ArrayLike, x0: ArrayLike = (), w: Optional[ArrayLike] = None,
116
+ wl: Optional[ArrayLike] = None, wr: Optional[ArrayLike] = None) -> None:
117
+ y, x0 = _array("y", y), _array("x0", x0)
118
+ n = y.shape[0] - 1
119
+ if n < 0:
120
+ raise ValueError("y must contain at least one level")
121
+ if x0.shape[0] != n:
122
+ raise ValueError(f"x0 must have {n} values (one per step), got {x0.shape[0]}")
123
+ if n == 0 and w is None and wl is None and wr is None:
124
+ w = ()
125
+ symmetric, w, wl, wr = _widths(n, w, wl, wr)
126
+ for name, value in (("y", y), ("x0", x0), ("w", w), ("wl", wl), ("wr", wr),
127
+ ("symmetric", symmetric)):
128
+ object.__setattr__(self, name, value)
129
+
130
+ def evaluate(self, x: ArrayLike) -> np.ndarray:
131
+ """Return the profile on the coordinates ``x``."""
132
+ x = np.asarray(x, dtype=np.float64)
133
+ pad = np.zeros(1)
134
+ return step_func(self.symmetric, x, self.y,
135
+ np.concatenate((pad, self.x0)), np.concatenate((pad, self.w)),
136
+ np.concatenate((pad, self.wl)), np.concatenate((pad, self.wr)))
137
+
138
+
139
+ def evaluate(profile: Union[None, float, Pulse, Steps], x: np.ndarray) -> np.ndarray:
140
+ """Evaluate a profile argument of a solver on the grid ``x``.
141
+
142
+ ``None`` gives zeros, a number a constant, a ``Pulse`` or ``Steps`` its values.
143
+ """
144
+ if profile is None:
145
+ return np.zeros(np.shape(x), dtype=np.float64)
146
+ if isinstance(profile, (Pulse, Steps)):
147
+ return profile.evaluate(x)
148
+ if isinstance(profile, (int, float, np.integer, np.floating)) and not isinstance(profile, bool):
149
+ return np.full(np.shape(x), float(profile), dtype=np.float64)
150
+ raise TypeError(f"expected None, a number, Pulse or Steps, got {type(profile).__name__}")
151
+
152
+
153
+ # -----------------------------------------------------------------------------
154
+ # numba kernels (arguments as in MB Solver's widgets)
155
+ # -----------------------------------------------------------------------------
156
+ @njit(nogil=True, cache=True)
157
+ def pulse_func(symmetric, x, a, x0, w, wl, wr):
158
+ """Sum of sech² components ``a_i sech²((x - x0_i) / w_i)`` on ``x``.
159
+
160
+ All of ``a, x0, w, wl, wr`` have length N; ``w`` is used when
161
+ ``symmetric``, otherwise ``wl`` left and ``wr`` right of ``x0_i``.
162
+ """
163
+ n = a.shape[0]
164
+ m = x.shape[0]
165
+ if x0.shape[0] != n or w.shape[0] != n or wl.shape[0] != n or wr.shape[0] != n:
166
+ raise ValueError("x0, w, wl and wr must have the same length as a")
167
+
168
+ y = np.zeros(m, dtype=np.float64)
169
+
170
+ for i in range(n):
171
+ ai = a[i]
172
+ if ai == 0.0:
173
+ continue
174
+
175
+ x0i = x0[i]
176
+ wi = w[i]
177
+ wli = wl[i]
178
+ wri = wr[i]
179
+
180
+ for j in range(m):
181
+ dx = x[j] - x0i
182
+ width = wi if symmetric else wli if dx < 0.0 else wri
183
+
184
+ if width == 0.0:
185
+ if dx == 0.0:
186
+ y[j] += ai
187
+ elif np.isinf(width):
188
+ y[j] += ai
189
+ else:
190
+ t = np.tanh(dx / width)
191
+ y[j] += ai * (1.0 - t * t)
192
+
193
+ return y
194
+
195
+
196
+ @njit(nogil=True, cache=True)
197
+ def step_func(symmetric, x, levels, x0, w, wl, wr):
198
+ """Smoothed steps between ``levels`` on ``x``.
199
+
200
+ ``levels`` has length N; ``x0, w, wl, wr`` also have length N with
201
+ index 0 unused (the layout of MB Solver's step widgets).
202
+ """
203
+ n = levels.shape[0]
204
+ m = x.shape[0]
205
+ if n == 0:
206
+ raise ValueError("levels must contain at least y^0")
207
+ if x0.shape[0] != n or w.shape[0] != n or wl.shape[0] != n or wr.shape[0] != n:
208
+ raise ValueError("x0, w, wl and wr must have the same length as levels")
209
+
210
+ y = np.empty(m, dtype=np.float64)
211
+ y[:] = levels[0]
212
+
213
+ for i in range(1, n):
214
+ di = levels[i] - levels[i - 1]
215
+ if di == 0.0:
216
+ continue
217
+
218
+ x0i = x0[i]
219
+ wi = w[i]
220
+ wli = wl[i]
221
+ wri = wr[i]
222
+
223
+ for j in range(m):
224
+ dx = x[j] - x0i
225
+ width = wi if symmetric else wli if dx < 0.0 else wri
226
+
227
+ if width == 0.0:
228
+ if dx > 0.0:
229
+ y[j] += di
230
+ elif dx == 0.0:
231
+ y[j] += 0.5 * di
232
+ elif np.isinf(width):
233
+ y[j] += 0.5 * di
234
+ else:
235
+ y[j] += 0.5 * di * (1.0 + np.tanh(dx / width))
236
+
237
+ return y
@@ -0,0 +1,131 @@
1
+ """Two-level Maxwell-Bloch model.
2
+
3
+ Dimensionless equations for population inversion ``w``, coherence ``R`` and
4
+ field ``A`` along a 1-D medium ``0 <= z <= 1``::
5
+
6
+ dw/dt = A R - (w - w0) / T1 + Λ(t)
7
+ dR/dt = -A w - (R - R0) / T2
8
+ dA/dz = -η R
9
+
10
+ with the field ``A(t, z=0)`` given at the entrance.
11
+ """
12
+
13
+ from __future__ import annotations
14
+
15
+ import math
16
+ from dataclasses import dataclass
17
+ from typing import Optional, Sequence, Union
18
+
19
+ import numpy as np
20
+
21
+ from mbsolver.grids import grids_1d
22
+ from mbsolver.profiles import Pulse, Steps, evaluate
23
+ from mbsolver.two_level._kernel_1d import solve_kernel
24
+
25
+ __all__ = ["Result", "solve_1d"]
26
+
27
+
28
+ @dataclass(frozen=True)
29
+ class Result:
30
+ """Output of :func:`solve_1d`.
31
+
32
+ Attributes:
33
+ time: Time grid, shape ``(nt,)``.
34
+ z: Positions of the stored planes, shape ``(n_planes,)``, from 0 to 1.
35
+ intensity: ``A²`` on the stored planes, shape ``(n_planes, nt)``.
36
+ w: Population inversion on the stored planes, shape ``(n_planes, nt)``.
37
+ pump: Total pump ``Λ(t)``.
38
+ field_in: Boundary field ``A(t, z=0)``.
39
+ t1, t2: Relaxation times on the time grid.
40
+ """
41
+
42
+ time: np.ndarray
43
+ z: np.ndarray
44
+ intensity: np.ndarray
45
+ w: np.ndarray
46
+ pump: np.ndarray
47
+ field_in: np.ndarray
48
+ t1: np.ndarray
49
+ t2: np.ndarray
50
+
51
+
52
+ Profile = Union[None, float, Pulse, Steps]
53
+
54
+
55
+ def solve_1d(
56
+ *,
57
+ t_max: float,
58
+ nt: int,
59
+ nz: int,
60
+ eta: float,
61
+ t1: Union[float, Steps],
62
+ t2: Union[float, Steps],
63
+ theta0: Optional[float] = None,
64
+ w0: Optional[float] = None,
65
+ R0: Optional[float] = None,
66
+ pump: Union[Profile, Sequence[Profile]] = None,
67
+ field_in: Profile = None,
68
+ n_planes: int = 2,
69
+ ) -> Result:
70
+ """Solve the two-level Maxwell-Bloch equations in one dimension.
71
+
72
+ Args:
73
+ t_max: Duration; the time grid runs from 0 to ``t_max``.
74
+ nt: Number of time points (>= 2).
75
+ nz: Number of z points (>= 2).
76
+ eta: Coupling coefficient ``η``.
77
+ t1: Population relaxation time, a number or :class:`Steps` (> 0; ``inf`` = none).
78
+ t2: Coherence relaxation time, a number or :class:`Steps` (> 0; ``inf`` = none).
79
+ theta0: Initial Bloch angle: ``w0 = cos θ0``, ``R0 = sin θ0``.
80
+ w0, R0: Initial population and coherence, instead of ``theta0``.
81
+ The medium also relaxes towards these values.
82
+ pump: Pump ``Λ(t)``: a :class:`Pulse`, a number, or a list of them (summed).
83
+ field_in: Field ``A(t)`` entering at ``z = 0``: a :class:`Pulse`, a number or ``None``.
84
+ n_planes: Number of evenly spaced z planes to return, from 0 to 1 (<= ``nz``).
85
+
86
+ Returns:
87
+ A :class:`Result`.
88
+
89
+ Example:
90
+ >>> res = solve_1d(t_max=200, nt=4000, nz=200, eta=5.0, theta0=0.01,
91
+ ... t1=1.0, t2=0.5, pump=Pulse(a=0.1, x0=50, w=10))
92
+ >>> exit_intensity = res.intensity[-1]
93
+ """
94
+ nt, nz, n_planes = int(nt), int(nz), int(n_planes)
95
+ if nt < 2 or nz < 2:
96
+ raise ValueError("nt and nz must be at least 2")
97
+ if not 1 <= n_planes <= nz:
98
+ raise ValueError(f"n_planes must be between 1 and nz ({nz}), got {n_planes}")
99
+ if not t_max > 0:
100
+ raise ValueError("t_max must be positive")
101
+
102
+ if theta0 is not None:
103
+ if w0 is not None or R0 is not None:
104
+ raise ValueError("give either theta0, or w0 and R0, not both")
105
+ w_init, R_init = math.cos(theta0), math.sin(theta0)
106
+ elif w0 is not None and R0 is not None:
107
+ w_init, R_init = float(w0), float(R0)
108
+ else:
109
+ raise ValueError("give the initial state as theta0, or as w0 and R0")
110
+
111
+ t, dt, z, dz = grids_1d(float(t_max), nt, nz)
112
+
113
+ pumps = pump if isinstance(pump, (list, tuple)) else [pump]
114
+ pump_t = np.zeros_like(t)
115
+ for p in pumps:
116
+ pump_t = pump_t + evaluate(p, t)
117
+ field_t = evaluate(field_in, t)
118
+ t1_t = evaluate(t1, t)
119
+ t2_t = evaluate(t2, t)
120
+ for name, arr in (("t1", t1_t), ("t2", t2_t)):
121
+ if np.any(arr <= 0):
122
+ raise ValueError(f"{name} must be positive")
123
+
124
+ w_z = np.full_like(z, w_init)
125
+ R_z = np.full_like(z, R_init)
126
+ w, intensity = solve_kernel(t, dt, z, dz, n_planes, w_z, R_z, field_t, pump_t,
127
+ t1_t, t2_t, float(eta))
128
+
129
+ planes = z[np.rint(np.linspace(0.0, 1.0, n_planes) * (nz - 1)).astype(np.int64)]
130
+ return Result(time=t, z=planes, intensity=intensity, w=w, pump=pump_t,
131
+ field_in=field_t, t1=t1_t, t2=t2_t)
@@ -0,0 +1,138 @@
1
+ """Two-level Maxwell-Bloch equations in one dimension (numba kernel).
2
+
3
+ Moved unchanged from MB Solver 8.13 (``solver/maxwell_bloch.py``):
4
+ RK4 in time for the population and coherence, explicit step in z for the field.
5
+
6
+ dw/dt = A R - (w - w0) / T1 + Λ
7
+ dR/dt = -A w - (R - R0) / T2
8
+ dA/dz = -η R
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ import numpy as np
14
+ from numba import njit
15
+
16
+
17
+ @njit(inline="always", nogil=True, cache=True)
18
+ def _dw_dt(w, R, A, kw, kR, w0, inv_t1, lambda_n):
19
+ """Population derivative for one RK stage."""
20
+ return A * (R + kR) - (w + kw - w0) * inv_t1 + lambda_n
21
+
22
+
23
+ @njit(inline="always", nogil=True, cache=True)
24
+ def _dR_dt(w, R, A, kw, kR, R0, inv_t2):
25
+ """Coherence derivative for one RK stage."""
26
+ return -A * (w + kw) - (R + kR - R0) * inv_t2
27
+
28
+
29
+ @njit(inline="always", nogil=True, cache=True)
30
+ def _dA_dz(eta, R, dz):
31
+ """Field increment across one spatial step."""
32
+ return -eta * R * dz
33
+
34
+
35
+ @njit(inline="always", nogil=True, cache=True)
36
+ def _rk4_dn_dR(w, R, A, inv_t1, inv_t2, w0, R0, lambda_n, dt, half_dt, dt6):
37
+ """One RK4 step for population and coherence; returns ``(dw, dR)``."""
38
+ kw_1 = _dw_dt(w, R, A, 0, 0, w0, inv_t1, lambda_n)
39
+ kR_1 = _dR_dt(w, R, A, 0, 0, R0, inv_t2)
40
+
41
+ kw_2 = _dw_dt(w, R, A, half_dt * kw_1, half_dt * kR_1, w0, inv_t1, lambda_n)
42
+ kR_2 = _dR_dt(w, R, A, half_dt * kw_1, half_dt * kR_1, R0, inv_t2)
43
+
44
+ kw_3 = _dw_dt(w, R, A, half_dt * kw_2, half_dt * kR_2, w0, inv_t1, lambda_n)
45
+ kR_3 = _dR_dt(w, R, A, half_dt * kw_2, half_dt * kR_2, R0, inv_t2)
46
+
47
+ kw_4 = _dw_dt(w, R, A, dt * kw_3, dt * kR_3, w0, inv_t1, lambda_n)
48
+ kR_4 = _dR_dt(w, R, A, dt * kw_3, dt * kR_3, R0, inv_t2)
49
+
50
+ dw = dt6 * (kw_1 + 2.0 * kw_2 + 2.0 * kw_3 + kw_4)
51
+ dR = dt6 * (kR_1 + 2.0 * kR_2 + 2.0 * kR_3 + kR_4)
52
+
53
+ return dw, dR
54
+
55
+
56
+ @njit(nogil=True, cache=True)
57
+ def solve_kernel(t, dt, z, dz, n_z_planes, w0, R0, A0, lambda_n, t1, t2, eta):
58
+ """Advance the system over the grid.
59
+
60
+ Args:
61
+ t, dt: Time grid and step.
62
+ z, dz: Spatial grid and step.
63
+ n_z_planes: Number of evenly spaced z planes to store (<= len(z)).
64
+ w0, R0: Initial (and equilibrium) population and coherence along z.
65
+ A0: Boundary field at z = 0 on the time grid.
66
+ lambda_n: Pump on the time grid.
67
+ t1, t2: Relaxation times on the time grid (> 0, may be inf).
68
+ eta: Coupling coefficient.
69
+
70
+ Returns:
71
+ ``(w, I)``, each of shape ``(n_z_planes, len(t))``.
72
+ """
73
+ nt = len(t)
74
+ nz = len(z)
75
+ z_planes = np.linspace(0.0, 1.0, n_z_planes)
76
+ z_plane_indices = np.rint(z_planes * (nz - 1)).astype(np.int64)
77
+
78
+ inv_t1 = 1.0 / t1
79
+ inv_t2 = 1.0 / t2
80
+ half_dt = 0.5 * dt
81
+ dt6 = dt / 6.0
82
+
83
+ w_old = np.empty(nt, dtype=np.float64)
84
+ R_old = np.empty(nt, dtype=np.float64)
85
+ A_old = np.empty(nt, dtype=np.float64)
86
+
87
+ w_new = np.empty_like(w_old)
88
+ R_new = np.empty_like(R_old)
89
+ A_new = np.empty_like(A_old)
90
+
91
+ w_result = np.empty((n_z_planes, nt), dtype=np.float64)
92
+ I_result = np.empty((n_z_planes, nt), dtype=np.float64)
93
+
94
+ w_old[0], R_old[0] = w0[0], R0[0]
95
+ A_old[0] = A0[0]
96
+ w0_0 = w0[0]
97
+ R0_0 = R0[0]
98
+
99
+ for t_idx in range(1, nt):
100
+ dw, dR = _rk4_dn_dR(
101
+ w_old[t_idx - 1], R_old[t_idx - 1], A_old[t_idx - 1],
102
+ inv_t1[t_idx - 1], inv_t2[t_idx - 1], w0_0, R0_0,
103
+ lambda_n[t_idx - 1], dt, half_dt, dt6,
104
+ )
105
+ w_old[t_idx] = w_old[t_idx - 1] + dw
106
+ R_old[t_idx] = R_old[t_idx - 1] + dR
107
+ A_old[t_idx] = A0[t_idx]
108
+
109
+ w_result[0, :] = w_old
110
+ I_result[0, :] = A_old * A_old
111
+ i = 1
112
+
113
+ for z_idx in range(1, nz):
114
+ w_new[0], R_new[0] = w0[z_idx], R0[z_idx]
115
+ A_new[0] = A_old[0] + _dA_dz(eta, R_old[0], dz)
116
+ w0_i = w0[z_idx]
117
+ R0_i = R0[z_idx]
118
+
119
+ for t_idx in range(1, nt):
120
+ dw, dR = _rk4_dn_dR(
121
+ w_new[t_idx - 1], R_new[t_idx - 1], A_new[t_idx - 1],
122
+ inv_t1[t_idx - 1], inv_t2[t_idx - 1], w0_i, R0_i,
123
+ lambda_n[t_idx - 1], dt, half_dt, dt6,
124
+ )
125
+ w_new[t_idx] = w_new[t_idx - 1] + dw
126
+ R_new[t_idx] = R_new[t_idx - 1] + dR
127
+ A_new[t_idx] = A_old[t_idx] + _dA_dz(eta, R_old[t_idx], dz)
128
+
129
+ if i < n_z_planes and z_plane_indices[i] == z_idx:
130
+ w_result[i, :] = w_new
131
+ I_result[i, :] = A_new * A_new
132
+ i += 1
133
+
134
+ w_old, w_new = w_new, w_old
135
+ R_old, R_new = R_new, R_old
136
+ A_old, A_new = A_new, A_old
137
+
138
+ return w_result, I_result
@@ -0,0 +1,127 @@
1
+ Metadata-Version: 2.4
2
+ Name: mbsolver
3
+ Version: 0.2.0
4
+ Summary: Fast Maxwell-Bloch solvers for masers, superradiance and coherent pulse propagation
5
+ Author-email: Vahid Anari <vahid.anari8@gmail.com>
6
+ License-Expression: MIT
7
+ Project-URL: Homepage, https://github.com/vahid-anari/mbsolver
8
+ Project-URL: Repository, https://github.com/vahid-anari/mbsolver
9
+ Project-URL: Issues, https://github.com/vahid-anari/mbsolver/issues
10
+ Keywords: maxwell-bloch,maser,superradiance,physics,astrophysics,numba,simulation
11
+ Classifier: Development Status :: 3 - Alpha
12
+ Classifier: Intended Audience :: Science/Research
13
+ Classifier: Operating System :: OS Independent
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Programming Language :: Python :: 3.13
19
+ Classifier: Topic :: Scientific/Engineering :: Physics
20
+ Classifier: Topic :: Scientific/Engineering :: Astronomy
21
+ Requires-Python: >=3.10
22
+ Description-Content-Type: text/markdown
23
+ License-File: LICENSE
24
+ Requires-Dist: numpy>=1.24
25
+ Requires-Dist: numba>=0.58
26
+ Provides-Extra: dev
27
+ Requires-Dist: build; extra == "dev"
28
+ Requires-Dist: twine; extra == "dev"
29
+ Requires-Dist: pytest; extra == "dev"
30
+ Dynamic: license-file
31
+
32
+ # mbsolver
33
+
34
+ Fast Maxwell-Bloch solvers in Python (numpy + numba) for astrophysical masers,
35
+ superradiance and coherent pulse propagation. Use them from scripts,
36
+ notebooks or cluster jobs; the
37
+ [Maxwell Bloch Studio](https://github.com/vahid-anari/maxwell-bloch-studio)
38
+ desktop app runs on the same code.
39
+
40
+ ```
41
+ pip install mbsolver
42
+ ```
43
+
44
+ ## Two-level model in one dimension
45
+
46
+ Dimensionless equations for the population inversion `w`, the coherence `R`
47
+ and the field `A` along a medium `0 ≤ z ≤ 1`:
48
+
49
+ ```
50
+ dw/dt = A R − (w − w0) / T1 + Λ(t)
51
+ dR/dt = −A w − (R − R0) / T2
52
+ dA/dz = −η R
53
+ ```
54
+
55
+ The field `A(t)` entering at `z = 0` is given; population and coherence are
56
+ advanced with fourth-order Runge-Kutta in time, the field step by step in `z`.
57
+
58
+ ```python
59
+ import matplotlib.pyplot as plt
60
+ from mbsolver import Pulse, Steps
61
+ from mbsolver.two_level import solve_1d
62
+
63
+ res = solve_1d(
64
+ t_max=300, nt=3001, nz=121, eta=40.0,
65
+ theta0=0.003, # or w0=..., R0=...
66
+ pump=[Pulse(a=[2e-3, 1e-3], x0=[60, 180], w=[15, 25])],
67
+ field_in=Pulse(a=1e-3, x0=30, w=5), # A(t) at z = 0
68
+ t1=Steps(y=[80, 40, 60], x0=[100, 200], w=[10, 0]),
69
+ t2=5.0,
70
+ n_planes=7, # stored z planes, 0 ... 1
71
+ )
72
+
73
+ plt.plot(res.time, res.intensity[-1]) # intensity leaving the medium
74
+ plt.show()
75
+ ```
76
+
77
+ ### Arguments
78
+
79
+ | Argument | Meaning |
80
+ |----------|---------|
81
+ | `t_max`, `nt` | time grid from 0 to `t_max` with `nt` points |
82
+ | `nz` | number of z points |
83
+ | `eta` | coupling coefficient η |
84
+ | `theta0`, or `w0` and `R0` | initial state (`w0 = cos θ0`, `R0 = sin θ0`); the medium also relaxes towards it |
85
+ | `t1`, `t2` | relaxation times: a number or `Steps` (> 0, `math.inf` for none) |
86
+ | `pump` | Λ(t): a `Pulse`, a number, or a list of them (summed) |
87
+ | `field_in` | A(t) at z = 0: a `Pulse`, a number, or `None` |
88
+ | `n_planes` | number of evenly spaced z planes returned (default 2: entrance and exit) |
89
+
90
+ ### Result
91
+
92
+ `res.time`, `res.z` (positions of the stored planes), `res.intensity` and
93
+ `res.w` (shape `(n_planes, nt)`), `res.pump`, `res.field_in`, `res.t1`, `res.t2`.
94
+
95
+ ## Profiles
96
+
97
+ `Pulse(a, x0, w)` is a sum of sech² pulses, `a_i sech²((t − x0_i) / w_i)`.
98
+ Give `wl` and `wr` instead of `w` for a different width on each side.
99
+
100
+ `Steps(y, x0, w)` goes from level `y[0]` to `y[1]` around `x0[0]`, then to
101
+ `y[2]` around `x0[1]`, and so on, each step smoothed as `(1 + tanh)/2` with
102
+ width `w` (0 for a sharp step). `x0` and `w` have one value fewer than `y`.
103
+
104
+ A zero width gives a single point (pulse) or a sharp step; an infinite width
105
+ a constant. `profile.evaluate(t)` returns the values on any grid.
106
+
107
+ ## Tests
108
+
109
+ ```
110
+ pip install -e ".[dev]"
111
+ pytest
112
+ ```
113
+
114
+ The tests include a reference run that must match the solver of the Maxwell
115
+ Bloch app (version 8.13) exactly, and physics checks (Bloch rotation at the
116
+ entrance, relaxation balance, uncoupled propagation).
117
+
118
+ ## Earlier name
119
+
120
+ Up to 0.1.0 this package was called `maxbloch`; 0.2.0 is the same code under
121
+ the new name. `maxbloch` stays on PyPI unchanged (0.1.1 only adds a
122
+ deprecation warning) but gets no updates; change `maxbloch` to `mbsolver` in
123
+ your imports.
124
+
125
+ ## License
126
+
127
+ MIT
@@ -0,0 +1,15 @@
1
+ LICENSE
2
+ README.md
3
+ pyproject.toml
4
+ src/mbsolver/__init__.py
5
+ src/mbsolver/grids.py
6
+ src/mbsolver/profiles.py
7
+ src/mbsolver.egg-info/PKG-INFO
8
+ src/mbsolver.egg-info/SOURCES.txt
9
+ src/mbsolver.egg-info/dependency_links.txt
10
+ src/mbsolver.egg-info/requires.txt
11
+ src/mbsolver.egg-info/top_level.txt
12
+ src/mbsolver/two_level/__init__.py
13
+ src/mbsolver/two_level/_kernel_1d.py
14
+ tests/test_profiles.py
15
+ tests/test_two_level.py
@@ -0,0 +1,7 @@
1
+ numpy>=1.24
2
+ numba>=0.58
3
+
4
+ [dev]
5
+ build
6
+ twine
7
+ pytest
@@ -0,0 +1,2 @@
1
+ maxbloch
2
+ mbsolver
@@ -0,0 +1,58 @@
1
+ """Tests for Pulse and Steps."""
2
+
3
+ import math
4
+
5
+ import numpy as np
6
+ import pytest
7
+
8
+ from mbsolver import Pulse, Steps
9
+
10
+ X = np.linspace(-10, 10, 201)
11
+
12
+
13
+ def test_pulse_symmetric():
14
+ y = Pulse(a=2.0, x0=1.0, w=3.0).evaluate(X)
15
+ np.testing.assert_allclose(y, 2.0 / np.cosh((X - 1.0) / 3.0) ** 2)
16
+
17
+
18
+ def test_pulse_sum_and_asymmetric():
19
+ p = Pulse(a=[1.0, 0.5], x0=[-2.0, 4.0], wl=[1.0, 2.0], wr=[3.0, 0.5])
20
+ left = X < -2.0
21
+ y = p.evaluate(X)
22
+ expected_first = np.where(left, 1 / np.cosh((X + 2) / 1.0) ** 2, 1 / np.cosh((X + 2) / 3.0) ** 2)
23
+ expected_second = 0.5 * np.where(X < 4.0, 1 / np.cosh((X - 4) / 2.0) ** 2,
24
+ 1 / np.cosh((X - 4) / 0.5) ** 2)
25
+ np.testing.assert_allclose(y, expected_first + expected_second)
26
+
27
+
28
+ def test_pulse_limits():
29
+ assert np.all(Pulse(a=3.0, x0=0.0, w=math.inf).evaluate(X) == 3.0)
30
+ y = Pulse(a=3.0, x0=0.0, w=0.0).evaluate(X)
31
+ assert y[100] == 3.0 and np.count_nonzero(y) == 1
32
+
33
+
34
+ def test_steps():
35
+ s = Steps(y=[1.0, 3.0, 2.0], x0=[-5.0, 5.0], w=[0.0, 2.0])
36
+ y = s.evaluate(X)
37
+ assert y[0] == pytest.approx(1.0, abs=1e-6)
38
+ assert y[50] == pytest.approx(2.0, abs=1e-3) # sharp step at x = -5 is halfway
39
+ assert y[80] == pytest.approx(3.0, abs=1e-3)
40
+ assert y[150] == pytest.approx(2.5) # smooth step at x = 5 is halfway
41
+ assert y[-1] == pytest.approx(2.0 + 0.5 * (1 - math.tanh(2.5)))
42
+
43
+
44
+ def test_single_level_is_constant():
45
+ assert np.all(Steps(y=4.0).evaluate(X) == 4.0)
46
+
47
+
48
+ @pytest.mark.parametrize("make, message", [
49
+ (lambda: Pulse(a=[1, 2], x0=[0], w=[1, 1]), "x0 must have 2"),
50
+ (lambda: Pulse(a=1, x0=0), "either w"),
51
+ (lambda: Pulse(a=1, x0=0, w=1, wl=1, wr=1), "not both"),
52
+ (lambda: Pulse(a=1, x0=0, w=-1), "negative"),
53
+ (lambda: Steps(y=[1, 2], x0=[0, 1], w=[1, 1]), "x0 must have 1"),
54
+ (lambda: Steps(y=[]), "at least one level"),
55
+ ])
56
+ def test_invalid(make, message):
57
+ with pytest.raises(ValueError, match=message):
58
+ make()
@@ -0,0 +1,92 @@
1
+ """Tests for mbsolver.two_level.solve_1d."""
2
+
3
+ import json
4
+ import math
5
+ from pathlib import Path
6
+
7
+ import numpy as np
8
+ import pytest
9
+
10
+ from mbsolver import Pulse, Steps
11
+ from mbsolver.two_level import solve_1d
12
+
13
+ DATA = Path(__file__).parent / "data"
14
+
15
+
16
+ def _pulse(p):
17
+ """MB Solver widget value -> Pulse."""
18
+ if p["symmetric"]:
19
+ return Pulse(a=p["a"], x0=p["x0"], w=p["w"])
20
+ return Pulse(a=p["a"], x0=p["x0"], wl=p["wl"], wr=p["wr"])
21
+
22
+
23
+ def _steps(s):
24
+ """MB Solver widget value -> number or Steps (index 0 of x0/w unused)."""
25
+ if not isinstance(s, dict):
26
+ return s
27
+ if s["symmetric"]:
28
+ return Steps(y=s["y"], x0=s["x0"][1:], w=s["w"][1:])
29
+ return Steps(y=s["y"], x0=s["x0"][1:], wl=s["wl"][1:], wr=s["wr"][1:])
30
+
31
+
32
+ def test_matches_mb_solver_8_13():
33
+ """Same results as MB Solver 8.13's own solver for a reference case."""
34
+ p = json.loads((DATA / "reference_params.json").read_text())
35
+ ref = np.load(DATA / "reference_two_level.npz")
36
+ res = solve_1d(
37
+ t_max=p["solve.grid.t_max"], nt=p["solve.grid.nt"], nz=p["solve.grid.nz"],
38
+ eta=p["solve.sample"]["eta"], theta0=p["solve.sample"]["theta0"],
39
+ pump=[_pulse(p[f"solve.pump.pulse{i}"]) for i in (1, 2, 3)],
40
+ field_in=_pulse(p["solve.bcs"]),
41
+ t1=_steps(p["solve.t1"]), t2=_steps(p["solve.t2"]),
42
+ n_planes=min(p["slice.z"]["arr_length"], p["solve.grid.nz"]),
43
+ )
44
+ for ours, theirs in (("time", "time"), ("intensity", "intensity"), ("w", "w"),
45
+ ("pump", "lambda_n"), ("field_in", "A0"), ("t1", "t1"), ("t2", "t2")):
46
+ np.testing.assert_allclose(getattr(res, ours), ref[theirs], rtol=1e-12, atol=1e-14,
47
+ err_msg=ours)
48
+
49
+ def test_entrance_is_a_bloch_rotation():
50
+ """At z = 0 without relaxation or pump, the field rotates the Bloch angle by minus its area."""
51
+ theta0, a, x0, width = 0.2, 0.05, 40.0, 8.0
52
+ res = solve_1d(t_max=100, nt=4001, nz=3, eta=1.0, theta0=theta0, t1=math.inf, t2=math.inf,
53
+ field_in=Pulse(a=a, x0=x0, w=width))
54
+ dt = res.time[1] - res.time[0]
55
+ area = np.concatenate(([0.0], np.cumsum(res.field_in[:-1]) * dt))
56
+ np.testing.assert_allclose(res.w[0], np.cos(theta0 - area), atol=1e-9)
57
+
58
+
59
+ def test_no_coupling_leaves_field_unchanged():
60
+ """With eta = 0 the field crosses the medium unchanged."""
61
+ res = solve_1d(t_max=50, nt=501, nz=11, eta=0.0, theta0=0.1, t1=10, t2=5,
62
+ field_in=Pulse(a=0.1, x0=20, w=5), n_planes=4)
63
+ for plane in res.intensity:
64
+ np.testing.assert_allclose(plane, res.field_in ** 2, rtol=0, atol=1e-15)
65
+
66
+
67
+ def test_relaxation_and_pump_balance():
68
+ """Without a field, w relaxes to w0 + T1 * pump for a constant pump."""
69
+ res = solve_1d(t_max=200, nt=2001, nz=2, eta=1.0, w0=0.0, R0=0.0, t1=10.0, t2=5.0,
70
+ pump=0.01)
71
+ assert res.w[-1, -1] == pytest.approx(0.1, rel=1e-6)
72
+
73
+
74
+ def test_result_shapes_and_planes():
75
+ res = solve_1d(t_max=10, nt=101, nz=21, eta=1.0, theta0=0.0, t1=1, t2=1, n_planes=5)
76
+ assert res.intensity.shape == res.w.shape == (5, 101)
77
+ np.testing.assert_allclose(res.z, [0, 0.25, 0.5, 0.75, 1.0])
78
+
79
+
80
+ @pytest.mark.parametrize("kwargs, message", [
81
+ (dict(nt=1), "nt and nz"),
82
+ (dict(n_planes=50), "n_planes"),
83
+ (dict(t1=0.0), "t1 must be positive"),
84
+ (dict(theta0=0.1, w0=1.0, R0=0.0), "not both"),
85
+ (dict(theta0=None), "initial state"),
86
+ (dict(pump="x"), "expected None"),
87
+ ])
88
+ def test_invalid_input(kwargs, message):
89
+ args = dict(t_max=10, nt=11, nz=11, eta=1.0, theta0=0.0, t1=1.0, t2=1.0)
90
+ args.update(kwargs)
91
+ with pytest.raises((ValueError, TypeError), match=message):
92
+ solve_1d(**args)