From 0adca4bcd7a709436f0782b836c3718315aefa40 Mon Sep 17 00:00:00 2001 From: Deepak Date: Sun, 2 Jan 2022 19:35:00 -0600 Subject: [PATCH] Adds bounds and discretisaiton --- pdme/model/fixed_z_plane_model.py | 38 ++++++++++++++++++++++ tests/model/test_fixed_z_discretization.py | 14 ++++++++ 2 files changed, 52 insertions(+) create mode 100644 tests/model/test_fixed_z_discretization.py diff --git a/pdme/model/fixed_z_plane_model.py b/pdme/model/fixed_z_plane_model.py index d4221bb..729cecb 100644 --- a/pdme/model/fixed_z_plane_model.py +++ b/pdme/model/fixed_z_plane_model.py @@ -1,4 +1,7 @@ import numpy +from dataclasses import dataclass +from typing import Tuple + from pdme.model.model import Model from pdme.measurement import DotMeasurement @@ -73,3 +76,38 @@ class FixedZPlaneModel(Model): 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 + + Parameters + ---------- + model : FixedZPlaneModel + 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. + ''' + model: FixedZPlaneModel + num_x: int + num_y: int + + def __post_init__(self): + self.cell_count = self.num_x * self.num_y + 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]) -> Tuple: + xi, yi = index + + # For this model, a point is (pz, sx, sy, w). + # We want to keep pz and w bounded, and restrict sx and sy. + return ([-numpy.inf, xi * self.x_step + self.model.xmin, yi * self.y_step + self.model.ymin, -numpy.inf], [numpy.inf, (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_x, self.num_y)) # type:ignore diff --git a/tests/model/test_fixed_z_discretization.py b/tests/model/test_fixed_z_discretization.py new file mode 100644 index 0000000..62be16d --- /dev/null +++ b/tests/model/test_fixed_z_discretization.py @@ -0,0 +1,14 @@ +from pdme.model.fixed_z_plane_model import FixedZPlaneModel, FixedZPlaneDiscretisation +import numpy + + +def test_fixed_z_plane_model_cost_and_jac_single(): + model = FixedZPlaneModel(4, -10, 10, -10, 10, 1) + discretisation = FixedZPlaneDiscretisation(model, 2, 5) + # x: (-10, 0) and (0, 10) + # y: (-10, -6, -2, 2, 6, 10) + assert discretisation.cell_count == 10 + assert discretisation.x_step == 10 + assert discretisation.y_step == 4 + numpy.testing.assert_allclose(discretisation.bounds((0, 0)), ((-numpy.inf, -10, -10, -numpy.inf), (numpy.inf, 0, -6, numpy.inf))) + numpy.testing.assert_allclose(list(discretisation.all_indices()), list(numpy.ndindex((2, 5))))