feat!: single dipole still outputs collection now to make interface consistent

This commit is contained in:
2022-04-30 16:58:45 -05:00
parent 58441b8ddb
commit cb3c280464
3 changed files with 34 additions and 15 deletions

View File

@@ -68,7 +68,6 @@ class SingleDipoleFixedMagnitudeModel(DipoleModel):
def get_monte_carlo_dipole_inputs(
self, n: int, max_frequency: float, rng_to_use: numpy.random.Generator = None
) -> numpy.ndarray:
# psw
rng: numpy.random.Generator
if rng_to_use is None:
@@ -76,16 +75,17 @@ class SingleDipoleFixedMagnitudeModel(DipoleModel):
else:
rng = rng_to_use
theta = 2 * numpy.pi * rng.random(n)
phi = numpy.arccos(2 * rng.random(n) - 1)
shape = (n, 1)
theta = 2 * numpy.pi * rng.random(shape)
phi = numpy.arccos(2 * rng.random(shape) - 1)
px = self.pfixed * numpy.cos(theta) * numpy.sin(phi)
py = self.pfixed * numpy.sin(theta) * numpy.sin(phi)
pz = self.pfixed * numpy.cos(phi)
sx = rng.uniform(self.xmin, self.xmax, n)
sy = rng.uniform(self.ymin, self.ymax, n)
sz = rng.uniform(self.zmin, self.zmax, n)
sx = rng.uniform(self.xmin, self.xmax, shape)
sy = rng.uniform(self.ymin, self.ymax, shape)
sz = rng.uniform(self.zmin, self.zmax, shape)
w = rng.uniform(1, max_frequency, n)
w = rng.uniform(1, max_frequency, shape)
return numpy.array([px, py, pz, sx, sy, sz, w]).T
return numpy.stack([px, py, pz, sx, sy, sz, w], axis=-1)