feat!: Removes unused models to make refactoring a bit easier
This commit is contained in:
@@ -1,14 +1,7 @@
|
||||
from pdme.model.model import Model, Discretisation
|
||||
from pdme.model.fixed_z_plane_model import FixedZPlaneModel
|
||||
from pdme.model.unrestricted_model import UnrestrictedModel
|
||||
from pdme.model.fixed_dipole_model import FixedDipoleModel
|
||||
from pdme.model.model import Model
|
||||
from pdme.model.fixed_magnitude_model import FixedMagnitudeModel
|
||||
|
||||
__all__ = [
|
||||
"Model",
|
||||
"Discretisation",
|
||||
"FixedZPlaneModel",
|
||||
"UnrestrictedModel",
|
||||
"FixedDipoleModel",
|
||||
"FixedMagnitudeModel",
|
||||
]
|
||||
|
||||
@@ -1,184 +0,0 @@
|
||||
import numpy
|
||||
import numpy.random
|
||||
from dataclasses import dataclass
|
||||
from typing import Sequence, Tuple
|
||||
import scipy.optimize
|
||||
from pdme.model.model import Model, Discretisation
|
||||
from pdme.measurement import (
|
||||
DotMeasurement,
|
||||
OscillatingDipoleArrangement,
|
||||
OscillatingDipole,
|
||||
)
|
||||
|
||||
|
||||
class FixedDipoleModel(Model):
|
||||
"""
|
||||
Model of oscillating dipole with a fixed dipole moment.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
p : numpy.ndarray
|
||||
The fixed dipole moment.
|
||||
|
||||
n : int
|
||||
The number of dipoles to assume.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
xmin: float,
|
||||
xmax: float,
|
||||
ymin: float,
|
||||
ymax: float,
|
||||
zmin: float,
|
||||
zmax: float,
|
||||
p: numpy.ndarray,
|
||||
n: int,
|
||||
) -> None:
|
||||
self.xmin = xmin
|
||||
self.xmax = xmax
|
||||
self.ymin = ymin
|
||||
self.ymax = ymax
|
||||
self.zmin = zmin
|
||||
self.zmax = zmax
|
||||
self.p = p
|
||||
self._n = n
|
||||
self.rng = numpy.random.default_rng()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"FixedDipoleModel({self.xmin}, {self.xmax}, {self.ymin}, {self.ymax}, {self.zmin}, {self.zmax}, {self.p}, {self.n()})"
|
||||
|
||||
# TODO: this signature doesn't make sense.
|
||||
def get_dipoles(self, frequency: float) -> OscillatingDipoleArrangement:
|
||||
s_pts = numpy.array(
|
||||
(
|
||||
self.rng.uniform(self.xmin, self.xmax),
|
||||
self.rng.uniform(self.ymin, self.ymax),
|
||||
self.rng.uniform(self.zmin, self.zmax),
|
||||
)
|
||||
)
|
||||
return OscillatingDipoleArrangement(
|
||||
[OscillatingDipole(self.p, s_pts, frequency)]
|
||||
)
|
||||
|
||||
def solution_single_dipole(self, pt: numpy.ndarray) -> OscillatingDipole:
|
||||
# assume length is 4.
|
||||
s = pt[0:3]
|
||||
w = pt[3]
|
||||
return OscillatingDipole(self.p, s, w)
|
||||
|
||||
def point_length(self) -> int:
|
||||
"""
|
||||
Dipole is constrained magnitude, but free orientation.
|
||||
Six degrees of freedom: (sx, sy, sz, w).
|
||||
"""
|
||||
return 4
|
||||
|
||||
def n(self) -> int:
|
||||
return self._n
|
||||
|
||||
def v_for_point_at_dot(self, dot: DotMeasurement, pt: numpy.ndarray) -> float:
|
||||
s = pt[0:3]
|
||||
w = pt[3]
|
||||
|
||||
diff = dot.r - s
|
||||
alpha = self.p.dot(diff) / (numpy.linalg.norm(diff) ** 3)
|
||||
b = (1 / numpy.pi) * (w / (w**2 + dot.f**2))
|
||||
return alpha**2 * b
|
||||
|
||||
def jac_for_point_at_dot(
|
||||
self, dot: DotMeasurement, pt: numpy.ndarray
|
||||
) -> numpy.ndarray:
|
||||
s = pt[0:3]
|
||||
w = pt[3]
|
||||
|
||||
diff = dot.r - s
|
||||
alpha = self.p.dot(diff) / (numpy.linalg.norm(diff) ** 3)
|
||||
b = (1 / numpy.pi) * (w / (w**2 + dot.f**2))
|
||||
|
||||
r_divs = (
|
||||
(
|
||||
-self.p / (numpy.linalg.norm(diff) ** 3)
|
||||
+ 3 * self.p.dot(diff) * diff / (numpy.linalg.norm(diff) ** 5)
|
||||
)
|
||||
* 2
|
||||
* alpha
|
||||
* b
|
||||
)
|
||||
|
||||
f2 = dot.f**2
|
||||
w2 = w**2
|
||||
|
||||
w_div = alpha**2 * (1 / numpy.pi) * ((f2 - w2) / ((f2 + w2) ** 2))
|
||||
|
||||
return numpy.concatenate((r_divs, w_div), axis=None)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FixedDipoleDiscretisation(Discretisation):
|
||||
"""
|
||||
Representation of a discretisation of a FixedDipoleDiscretisation.
|
||||
Also captures a rough maximum value of dipole.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model : FixedDipoleModel
|
||||
The parent model of the discretisation.
|
||||
|
||||
num_x : int
|
||||
The number of partitions of the x axis.
|
||||
|
||||
num_y : int
|
||||
The number of partitions of the y axis.
|
||||
|
||||
num_z : int
|
||||
The number of partitions of the z axis.
|
||||
"""
|
||||
|
||||
model: FixedDipoleModel
|
||||
num_x: int
|
||||
num_y: int
|
||||
num_z: int
|
||||
|
||||
def __post_init__(self):
|
||||
self.cell_count = self.num_x * self.num_y * self.num_z
|
||||
self.x_step = (self.model.xmax - self.model.xmin) / self.num_x
|
||||
self.y_step = (self.model.ymax - self.model.ymin) / self.num_y
|
||||
self.z_step = (self.model.zmax - self.model.zmin) / self.num_z
|
||||
|
||||
def bounds(self, index: Tuple[float, ...]) -> Tuple:
|
||||
xi, yi, zi = index
|
||||
|
||||
# For this model, a point is (sx, sx, sy, w).
|
||||
# We want to keep w unbounded, restrict sx, sy, sz, px and py based on step.
|
||||
return (
|
||||
[
|
||||
xi * self.x_step + self.model.xmin,
|
||||
yi * self.y_step + self.model.ymin,
|
||||
zi * self.z_step + self.model.zmin,
|
||||
-numpy.inf,
|
||||
],
|
||||
[
|
||||
(xi + 1) * self.x_step + self.model.xmin,
|
||||
(yi + 1) * self.y_step + self.model.ymin,
|
||||
(zi + 1) * self.z_step + self.model.zmin,
|
||||
numpy.inf,
|
||||
],
|
||||
)
|
||||
|
||||
def all_indices(self) -> numpy.ndindex:
|
||||
# see https://github.com/numpy/numpy/issues/20706 for why this is a mypy problem.
|
||||
return numpy.ndindex((self.num_x, self.num_y, self.num_z)) # type:ignore
|
||||
|
||||
def solve_for_index(
|
||||
self, dots: Sequence[DotMeasurement], index: Tuple[float, ...]
|
||||
) -> scipy.optimize.OptimizeResult:
|
||||
bounds = self.bounds(index)
|
||||
sx_mean = (bounds[0][0] + bounds[1][0]) / 2
|
||||
sy_mean = (bounds[0][1] + bounds[1][1]) / 2
|
||||
sz_mean = (bounds[0][2] + bounds[1][2]) / 2
|
||||
return self.model.solve(
|
||||
dots,
|
||||
initial_pt=numpy.array([sx_mean, sy_mean, sz_mean, 0.1]),
|
||||
bounds=bounds,
|
||||
)
|
||||
@@ -1,9 +1,6 @@
|
||||
import numpy
|
||||
import numpy.random
|
||||
from dataclasses import dataclass
|
||||
from typing import Sequence, Tuple
|
||||
import scipy.optimize
|
||||
from pdme.model.model import Model, Discretisation
|
||||
from pdme.model.model import Model
|
||||
from pdme.measurement import (
|
||||
DotMeasurement,
|
||||
OscillatingDipole,
|
||||
@@ -182,96 +179,3 @@ class FixedMagnitudeModel(Model):
|
||||
w_div = alpha**2 * (1 / numpy.pi) * ((f2 - w2) / ((f2 + w2) ** 2))
|
||||
|
||||
return numpy.concatenate((theta_div, phi_div, r_divs, w_div), axis=None)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FixedMagnitudeDiscretisation(Discretisation):
|
||||
"""
|
||||
Representation of a discretisation of a FixedMagnitudeDiscretisation.
|
||||
Also captures a rough maximum value of dipole.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model : FixedMagnitudeModel
|
||||
The parent model of the discretisation.
|
||||
|
||||
num_ptheta: int
|
||||
The number of partitions of ptheta.
|
||||
|
||||
num_pphi: int
|
||||
The number of partitions of pphi.
|
||||
|
||||
num_x : int
|
||||
The number of partitions of the x axis.
|
||||
|
||||
num_y : int
|
||||
The number of partitions of the y axis.
|
||||
|
||||
num_z : int
|
||||
The number of partitions of the z axis.
|
||||
"""
|
||||
|
||||
model: FixedMagnitudeModel
|
||||
num_ptheta: int
|
||||
num_pphi: int
|
||||
num_x: int
|
||||
num_y: int
|
||||
num_z: int
|
||||
|
||||
def __post_init__(self):
|
||||
self.cell_count = self.num_x * self.num_y * self.num_z
|
||||
self.x_step = (self.model.xmax - self.model.xmin) / self.num_x
|
||||
self.y_step = (self.model.ymax - self.model.ymin) / self.num_y
|
||||
self.z_step = (self.model.zmax - self.model.zmin) / self.num_z
|
||||
self.h_step = 2 / self.num_ptheta
|
||||
self.phi_step = 2 * numpy.pi / self.num_pphi
|
||||
|
||||
def bounds(self, index: Tuple[float, ...]) -> Tuple:
|
||||
pthetai, pphii, xi, yi, zi = index
|
||||
|
||||
# For this model, a point is (p_theta, p_phi, sx, sx, sy, w).
|
||||
# We want to keep w unbounded, restrict sx, sy, sz, px and py based on step.
|
||||
return (
|
||||
[
|
||||
numpy.arccos(1 - pthetai * self.h_step),
|
||||
pphii * self.phi_step,
|
||||
xi * self.x_step + self.model.xmin,
|
||||
yi * self.y_step + self.model.ymin,
|
||||
zi * self.z_step + self.model.zmin,
|
||||
-numpy.inf,
|
||||
],
|
||||
[
|
||||
numpy.arccos(1 - (pthetai + 1) * self.h_step),
|
||||
(pphii + 1) * self.phi_step,
|
||||
(xi + 1) * self.x_step + self.model.xmin,
|
||||
(yi + 1) * self.y_step + self.model.ymin,
|
||||
(zi + 1) * self.z_step + self.model.zmin,
|
||||
numpy.inf,
|
||||
],
|
||||
)
|
||||
|
||||
def get_model(self) -> Model:
|
||||
return self.model
|
||||
|
||||
def all_indices(self) -> numpy.ndindex:
|
||||
# see https://github.com/numpy/numpy/issues/20706 for why this is a mypy problem.
|
||||
return numpy.ndindex(
|
||||
(self.num_ptheta, self.num_pphi, self.num_x, self.num_y, self.num_z)
|
||||
) # type:ignore
|
||||
|
||||
def solve_for_index(
|
||||
self, dots: Sequence[DotMeasurement], index: Tuple[float, ...]
|
||||
) -> scipy.optimize.OptimizeResult:
|
||||
bounds = self.bounds(index)
|
||||
ptheta_mean = (bounds[0][0] + bounds[1][0]) / 2
|
||||
pphi_mean = (bounds[0][1] + bounds[1][1]) / 2
|
||||
sx_mean = (bounds[0][2] + bounds[1][2]) / 2
|
||||
sy_mean = (bounds[0][3] + bounds[1][3]) / 2
|
||||
sz_mean = (bounds[0][4] + bounds[1][4]) / 2
|
||||
return self.model.solve(
|
||||
dots,
|
||||
initial_pt=numpy.array(
|
||||
[ptheta_mean, pphi_mean, sx_mean, sy_mean, sz_mean, 0.1]
|
||||
),
|
||||
bounds=bounds,
|
||||
)
|
||||
|
||||
@@ -1,173 +0,0 @@
|
||||
import numpy
|
||||
from dataclasses import dataclass
|
||||
from typing import Tuple, Sequence
|
||||
import scipy.optimize
|
||||
|
||||
from pdme.model.model import Model
|
||||
from pdme.measurement import DotMeasurement
|
||||
|
||||
|
||||
class FixedZPlaneModel(Model):
|
||||
"""
|
||||
Model of oscillating dipoles constrained to lie within a plane.
|
||||
Additionally, each dipole is assumed to be orientated in the plus or minus z direction.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
z : float
|
||||
The z position of the plane where dipoles are constrained to lie.
|
||||
|
||||
xmin : float
|
||||
The minimum x value for dipoles.
|
||||
|
||||
xmax : float
|
||||
The maximum x value for dipoles.
|
||||
|
||||
ymin : float
|
||||
The minimum y value for dipoles.
|
||||
|
||||
ymax : float
|
||||
The maximum y value for dipoles.
|
||||
|
||||
n : int
|
||||
The number of dipoles to assume.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, z: float, xmin: float, xmax: float, ymin: float, ymax: float, n: int
|
||||
) -> None:
|
||||
self.z = z
|
||||
self.xmin = xmin
|
||||
self.xmax = xmax
|
||||
self.ymin = ymin
|
||||
self.ymax = ymax
|
||||
self._n = n
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"FixedZPlaneModel({self.z}, {self.xmin}, {self.xmax}, {self.ymin}, {self.ymax}, {self.n()})"
|
||||
|
||||
def point_length(self) -> int:
|
||||
"""
|
||||
Dipole is constrained in this model to have (px, py, pz) = (0, 0, pz) and (sx, sy, sz) = (sx, sy, self.z).
|
||||
With some frequency w, there are four degrees of freedom: (pz, sx, sy, w).
|
||||
"""
|
||||
return 4
|
||||
|
||||
def n(self) -> int:
|
||||
return self._n
|
||||
|
||||
def v_for_point_at_dot(self, dot: DotMeasurement, pt: numpy.ndarray) -> float:
|
||||
p = numpy.array([0, 0, pt[0]])
|
||||
s = numpy.array([pt[1], pt[2], self.z])
|
||||
w = pt[3]
|
||||
|
||||
diff = dot.r - s
|
||||
alpha = p.dot(diff) / (numpy.linalg.norm(diff) ** 3)
|
||||
b = (1 / numpy.pi) * (w / (w**2 + dot.f**2))
|
||||
return alpha**2 * b
|
||||
|
||||
def jac_for_point_at_dot(
|
||||
self, dot: DotMeasurement, pt: numpy.ndarray
|
||||
) -> numpy.ndarray:
|
||||
p = numpy.array([0, 0, pt[0]])
|
||||
s = numpy.array([pt[1], pt[2], self.z])
|
||||
w = pt[3]
|
||||
|
||||
diff = dot.r - s
|
||||
alpha = p.dot(diff) / (numpy.linalg.norm(diff) ** 3)
|
||||
b = (1 / numpy.pi) * (w / (w**2 + dot.f**2))
|
||||
|
||||
p_divs = (
|
||||
2 * alpha * diff[2] / (numpy.linalg.norm(diff) ** 3) * b
|
||||
) # only need the z component.
|
||||
|
||||
r_divs = (
|
||||
(
|
||||
-p[0:2] / (numpy.linalg.norm(diff) ** 3)
|
||||
+ 3 * p.dot(diff) * diff[0:2] / (numpy.linalg.norm(diff) ** 5)
|
||||
)
|
||||
* 2
|
||||
* alpha
|
||||
* b
|
||||
)
|
||||
|
||||
f2 = dot.f**2
|
||||
w2 = w**2
|
||||
|
||||
w_div = alpha**2 * (1 / numpy.pi) * ((f2 - w2) / ((f2 + w2) ** 2))
|
||||
|
||||
return numpy.concatenate((p_divs, r_divs, w_div), axis=None)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FixedZPlaneDiscretisation:
|
||||
"""
|
||||
Representation of a discretisation of a FixedZPlaneModel.
|
||||
Also captures a rough maximum value of dipole.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model : FixedZPlaneModel
|
||||
The parent model of the discretisation.
|
||||
|
||||
num_pz: int
|
||||
The number of partitions of pz.
|
||||
|
||||
num_x : int
|
||||
The number of partitions of the x axis.
|
||||
|
||||
num_y : int
|
||||
The number of partitions of the y axis.
|
||||
"""
|
||||
|
||||
model: FixedZPlaneModel
|
||||
num_pz: int
|
||||
num_x: int
|
||||
num_y: int
|
||||
max_pz: int
|
||||
|
||||
def __post_init__(self):
|
||||
self.cell_count = self.num_x * self.num_y
|
||||
self.pz_step = (2 * self.max_pz) / self.num_pz
|
||||
self.x_step = (self.model.xmax - self.model.xmin) / self.num_x
|
||||
self.y_step = (self.model.ymax - self.model.ymin) / self.num_y
|
||||
|
||||
def bounds(
|
||||
self, index: Tuple[float, float, float]
|
||||
) -> Tuple[numpy.ndarray, numpy.ndarray]:
|
||||
pzi, xi, yi = index
|
||||
|
||||
# For this model, a point is (pz, sx, sy, w).
|
||||
# We want to keep w bounded, and restrict pz, sx and sy based on step.
|
||||
return (
|
||||
numpy.array(
|
||||
(
|
||||
pzi * self.pz_step - self.max_pz,
|
||||
xi * self.x_step + self.model.xmin,
|
||||
yi * self.y_step + self.model.ymin,
|
||||
-numpy.inf,
|
||||
)
|
||||
),
|
||||
numpy.array(
|
||||
(
|
||||
(pzi + 1) * self.pz_step - self.max_pz,
|
||||
(xi + 1) * self.x_step + self.model.xmin,
|
||||
(yi + 1) * self.y_step + self.model.ymin,
|
||||
numpy.inf,
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
def all_indices(self) -> numpy.ndindex:
|
||||
# see https://github.com/numpy/numpy/issues/20706 for why this is a mypy problem.
|
||||
return numpy.ndindex((self.num_pz, self.num_x, self.num_y)) # type:ignore
|
||||
|
||||
def solve_for_index(
|
||||
self, dots: Sequence[DotMeasurement], index: Tuple[float, float, float]
|
||||
) -> scipy.optimize.OptimizeResult:
|
||||
bounds = self.bounds(index)
|
||||
pz_mean = (bounds[0][0] + bounds[1][0]) / 2
|
||||
sx_mean = (bounds[0][1] + bounds[1][1]) / 2
|
||||
sy_mean = (bounds[0][2] + bounds[1][2]) / 2
|
||||
# I don't care about the typing here at the moment.
|
||||
return self.model.solve(dots, initial_pt=numpy.array((pz_mean, sx_mean, sy_mean, 0.1)), bounds=bounds) # type: ignore
|
||||
@@ -1,7 +1,7 @@
|
||||
import numpy
|
||||
import numpy.random
|
||||
import scipy.optimize
|
||||
from typing import Callable, Sequence, Tuple, List
|
||||
from typing import Callable, Sequence, List
|
||||
from pdme.measurement import (
|
||||
DotMeasurement,
|
||||
OscillatingDipoleArrangement,
|
||||
@@ -136,19 +136,3 @@ class Model:
|
||||
result.x, self.point_length()
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
class Discretisation:
|
||||
def bounds(self, index: Tuple[float, ...]) -> Tuple:
|
||||
raise NotImplementedError
|
||||
|
||||
def all_indices(self) -> numpy.ndindex:
|
||||
raise NotImplementedError
|
||||
|
||||
def solve_for_index(
|
||||
self, dots: Sequence[DotMeasurement], index: Tuple
|
||||
) -> scipy.optimize.OptimizeResult:
|
||||
raise NotImplementedError
|
||||
|
||||
def get_model(self) -> Model:
|
||||
raise NotImplementedError
|
||||
|
||||
@@ -1,216 +0,0 @@
|
||||
import numpy
|
||||
from dataclasses import dataclass
|
||||
from typing import Sequence, Tuple
|
||||
import scipy.optimize
|
||||
from pdme.model.model import Model, Discretisation
|
||||
from pdme.measurement import (
|
||||
DotMeasurement,
|
||||
OscillatingDipoleArrangement,
|
||||
OscillatingDipole,
|
||||
)
|
||||
|
||||
|
||||
class UnrestrictedModel(Model):
|
||||
"""
|
||||
Model of oscillating dipoles with no restrictions.
|
||||
Additionally, each dipole is assumed to be orientated in the plus or minus z direction.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
n : int
|
||||
The number of dipoles to assume.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
xmin: float,
|
||||
xmax: float,
|
||||
ymin: float,
|
||||
ymax: float,
|
||||
zmin: float,
|
||||
zmax: float,
|
||||
max_p: float,
|
||||
n: int,
|
||||
) -> None:
|
||||
self.xmin = xmin
|
||||
self.xmax = xmax
|
||||
self.ymin = ymin
|
||||
self.ymax = ymax
|
||||
self.zmin = zmin
|
||||
self.zmax = zmax
|
||||
self.max_p = max_p
|
||||
self._n = n
|
||||
self.rng = numpy.random.default_rng()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"UnrestrictedModel({self.xmin}, {self.xmax}, {self.ymin}, {self.ymax}, {self.zmin}, {self.zmax}, {self.max_p}, {self.n()})"
|
||||
|
||||
def point_length(self) -> int:
|
||||
"""
|
||||
Dipole is unconstrained in this model.
|
||||
All seven degrees of freedom: (px, py, pz, sx, sy, sz, w).
|
||||
"""
|
||||
return 7
|
||||
|
||||
def n(self) -> int:
|
||||
return self._n
|
||||
|
||||
def v_for_point_at_dot(self, dot: DotMeasurement, pt: numpy.ndarray) -> float:
|
||||
p = pt[0:3]
|
||||
s = pt[3:6]
|
||||
w = pt[6]
|
||||
|
||||
diff = dot.r - s
|
||||
alpha = p.dot(diff) / (numpy.linalg.norm(diff) ** 3)
|
||||
b = (1 / numpy.pi) * (w / (w**2 + dot.f**2))
|
||||
return alpha**2 * b
|
||||
|
||||
def jac_for_point_at_dot(
|
||||
self, dot: DotMeasurement, pt: numpy.ndarray
|
||||
) -> numpy.ndarray:
|
||||
p = pt[0:3]
|
||||
s = pt[3:6]
|
||||
w = pt[6]
|
||||
|
||||
diff = dot.r - s
|
||||
alpha = p.dot(diff) / (numpy.linalg.norm(diff) ** 3)
|
||||
b = (1 / numpy.pi) * (w / (w**2 + dot.f**2))
|
||||
|
||||
p_divs = 2 * alpha * diff / (numpy.linalg.norm(diff) ** 3) * b
|
||||
|
||||
r_divs = (
|
||||
(
|
||||
-p / (numpy.linalg.norm(diff) ** 3)
|
||||
+ 3 * p.dot(diff) * diff / (numpy.linalg.norm(diff) ** 5)
|
||||
)
|
||||
* 2
|
||||
* alpha
|
||||
* b
|
||||
)
|
||||
|
||||
f2 = dot.f**2
|
||||
w2 = w**2
|
||||
|
||||
w_div = alpha**2 * (1 / numpy.pi) * ((f2 - w2) / ((f2 + w2) ** 2))
|
||||
|
||||
return numpy.concatenate((p_divs, r_divs, w_div), axis=None)
|
||||
|
||||
def get_dipoles(self, frequency: float) -> OscillatingDipoleArrangement:
|
||||
theta = numpy.arccos(self.rng.uniform(-1, 1))
|
||||
phi = self.rng.uniform(0, 2 * numpy.pi)
|
||||
p = self.rng.uniform(0, self.max_p)
|
||||
px = p * numpy.sin(theta) * numpy.cos(phi)
|
||||
py = p * numpy.sin(theta) * numpy.sin(phi)
|
||||
pz = p * numpy.cos(theta)
|
||||
s_pts = numpy.array(
|
||||
(
|
||||
self.rng.uniform(self.xmin, self.xmax),
|
||||
self.rng.uniform(self.ymin, self.ymax),
|
||||
self.rng.uniform(self.zmin, self.zmax),
|
||||
)
|
||||
)
|
||||
return OscillatingDipoleArrangement(
|
||||
[OscillatingDipole(numpy.array([px, py, pz]), s_pts, frequency)]
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class UnrestrictedDiscretisation(Discretisation):
|
||||
"""
|
||||
Representation of a discretisation of a UnrestrictedModel.
|
||||
Also captures a rough maximum value of dipole.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model : UnrestrictedModel
|
||||
The parent model of the discretisation.
|
||||
|
||||
num_px: int
|
||||
The number of partitions of the px.
|
||||
|
||||
num_py: int
|
||||
The number of partitions of the py.
|
||||
|
||||
num_pz: int
|
||||
The number of partitions of pz.
|
||||
|
||||
num_x : int
|
||||
The number of partitions of the x axis.
|
||||
|
||||
num_y : int
|
||||
The number of partitions of the y axis.
|
||||
|
||||
num_z : int
|
||||
The number of partitions of the z axis.
|
||||
|
||||
max_p : int
|
||||
The maximum p coordinate in any direction.
|
||||
"""
|
||||
|
||||
model: UnrestrictedModel
|
||||
num_px: int
|
||||
num_py: int
|
||||
num_pz: int
|
||||
num_x: int
|
||||
num_y: int
|
||||
num_z: int
|
||||
|
||||
def __post_init__(self):
|
||||
self.max_p = self.model.max_p
|
||||
self.cell_count = self.num_x * self.num_y * self.num_z
|
||||
self.x_step = (self.model.xmax - self.model.xmin) / self.num_x
|
||||
self.y_step = (self.model.ymax - self.model.ymin) / self.num_y
|
||||
self.z_step = (self.model.zmax - self.model.zmin) / self.num_z
|
||||
self.px_step = 2 * self.max_p / self.num_px
|
||||
self.py_step = 2 * self.max_p / self.num_py
|
||||
self.pz_step = 2 * self.max_p / self.num_pz
|
||||
|
||||
def bounds(self, index: Tuple[float, ...]) -> Tuple:
|
||||
pxi, pyi, pzi, xi, yi, zi = index
|
||||
|
||||
# For this model, a point is (px, py, pz, sx, sx, sy, w).
|
||||
# We want to keep w unbounded, restrict sx, sy, sz, px and py based on step.
|
||||
return (
|
||||
[
|
||||
pxi * self.px_step - self.max_p,
|
||||
pyi * self.py_step - self.max_p,
|
||||
pzi * self.pz_step - self.max_p,
|
||||
xi * self.x_step + self.model.xmin,
|
||||
yi * self.y_step + self.model.ymin,
|
||||
zi * self.z_step + self.model.zmin,
|
||||
-numpy.inf,
|
||||
],
|
||||
[
|
||||
(pxi + 1) * self.px_step - self.max_p,
|
||||
(pyi + 1) * self.py_step - self.max_p,
|
||||
(pzi + 1) * self.pz_step - self.max_p,
|
||||
(xi + 1) * self.x_step + self.model.xmin,
|
||||
(yi + 1) * self.y_step + self.model.ymin,
|
||||
(zi + 1) * self.z_step + self.model.zmin,
|
||||
numpy.inf,
|
||||
],
|
||||
)
|
||||
|
||||
def all_indices(self) -> numpy.ndindex:
|
||||
# see https://github.com/numpy/numpy/issues/20706 for why this is a mypy problem.
|
||||
return numpy.ndindex(
|
||||
(self.num_px, self.num_py, self.num_pz, self.num_x, self.num_y, self.num_z)
|
||||
) # type:ignore
|
||||
|
||||
def solve_for_index(
|
||||
self, dots: Sequence[DotMeasurement], index: Tuple[float, ...]
|
||||
) -> scipy.optimize.OptimizeResult:
|
||||
bounds = self.bounds(index)
|
||||
px_mean = (bounds[0][0] + bounds[1][0]) / 2
|
||||
py_mean = (bounds[0][1] + bounds[1][1]) / 2
|
||||
pz_mean = (bounds[0][2] + bounds[1][2]) / 2
|
||||
sx_mean = (bounds[0][3] + bounds[1][3]) / 2
|
||||
sy_mean = (bounds[0][4] + bounds[1][4]) / 2
|
||||
sz_mean = (bounds[0][5] + bounds[1][5]) / 2
|
||||
return self.model.solve(
|
||||
dots,
|
||||
initial_pt=numpy.array(
|
||||
[px_mean, py_mean, pz_mean, sx_mean, sy_mean, sz_mean, 0.1]
|
||||
),
|
||||
bounds=bounds,
|
||||
)
|
||||
Reference in New Issue
Block a user