fix: lets you get model from discretisation for fixed magnitude models in type hintable way

This commit is contained in:
Deepak Mallubhotla 2022-03-06 16:38:48 -06:00
parent 7a2f34de3b
commit 548dcfebfc
Signed by: deepak
GPG Key ID: BEBAEBF28083E022
2 changed files with 6 additions and 0 deletions

View File

@ -190,6 +190,9 @@ class FixedMagnitudeDiscretisation(Discretisation):
]
)
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

View File

@ -114,3 +114,6 @@ class Discretisation():
def solve_for_index(self, dots: Sequence[DotMeasurement], index: Tuple) -> scipy.optimize.OptimizeResult:
raise NotImplementedError
def get_model(self) -> Model:
raise NotImplementedError