def build_n_fields(self, N, samples=2500):
"""Builds a dataset with N fields and random masses
:param N: number of fields
:param samples: number of data points
:returns: input data with shape (n, field_size, field_size, 2) and output data (n, 1)
"""
field_potential = []
vector_fields = []
field = []
for i in range(N):
print(f"Produced {i + 1} fields\r", end="", flush=True)
field_potential.append(
fields.trap(
fields.scalar_field(size=self.sim_field_size,
rng=self.rng)))
vector_fields.append(fields.vector_field(field_potential[i]))
field.append(
fields.downsize_field(field_potential[i],
self.pred_field_shape))
print("")
masses = self.draw_mass(samples, self.rng)
x = []
ignored = 0
field_inds = np.zeros(samples)
for sample in range(samples):
print(f"Produced {sample+1} samples\r", end="", flush=True)
if N is 1:
j = 0
else:
j = np.random.randint(0, N)
try:
initial_pos = (self.sim_field_size // 2,
self.sim_field_size // 2)
positions = simulation.simulate(vector_fields[j],
masses[sample],
initial_pos,
steps=self.steps)
pos = simulation.four_warm_position(self.sim_field_shape,
self.pred_field_shape,
positions[-1])
x.append(np.stack([field[j], pos], axis=-1))
field_inds[sample] = j
except IndexError:
# Sometimes the particle achieves incredible acceleration and
# goes completely haywire
print('\nWARNING: Ignored simulation which gave IndexError')
ignored += 1
print("")
print(f"Ignored {ignored} simulations which gave IndexError")
print(f"Size of data set: {samples - ignored}")
return np.array(x), masses, field_inds
def _simulate(self, rng, mass, steps=500):
"""Simulate a particle in a random field
:param rng: Numpy random number generator to use
:param steps: Number of steps to simulate the particle for
:return: A tuple with the field potential, and the particle position after
``steps`` simulation steps.
"""
field_potential = fields.trap(
fields.scalar_field(size=self.sim_field_size, rng=rng))
vector_field = fields.vector_field(field_potential)
initial_pos = (self.sim_field_size // 2, self.sim_field_size // 2)
pos = simulation.simulate(vector_field, mass, initial_pos, steps=steps)
return field_potential, pos
def test_vector_field_shape():
vectors = fields.vector_field(np.zeros((10, 20)))
assert vectors.shape == (10, 20, 2)
def test_simulate():
"""Smoke test to ensure it doesn't error out."""
field = fields.scalar_field()
vectors = fields.vector_field(field)
simulation.simulate(vectors, 1, (12, 12))