def _build_batch(self, samples, rng):
skipped_steps = 10
x, y = [], []
ignored_samples = 0
for sample in range(samples):
valid_sample = False
while not valid_sample:
try:
mass, field, positions = self.simulate(rng)
field = np.repeat(fields.downsize_field(
field, self.pred_field_shape)[np.newaxis],
self.model.num_timesteps,
axis=0)
pos = np.zeros(
(self.model.num_timesteps, *self.pred_field_shape))
for i, position in enumerate(positions[::skipped_steps]):
pos[i] = simulation.four_warm_position(
self.sim_field_shape, self.pred_field_shape,
position)
x.append(np.stack([field, pos], axis=-1))
y.append([mass])
valid_sample = True
except IndexError:
# Sometimes the particle achieves incredible acceleration and
# goes completely haywire
ignored_samples += 1
return np.array(x), np.array(y), ignored_samples
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 build_alt_set(self, samples=5000, time_steps=10, sim_steps=200):
mass = self.draw_mass(samples, self.rng)
steps_skip = sim_steps // time_steps
x = np.zeros((samples, 32, 32, 2))
for sample in range(samples):
print(f"Produced {sample+1} samples\r", end="", flush=True)
field, positions = self._simulate(self.rng, mass[sample],
sim_steps)
field = fields.downsize_field(field, self.pred_field_shape)
positions = positions[0::steps_skip]
fw_positions = np.empty(
(time_steps, self.pred_field_size, self.pred_field_size))
for i, pos in enumerate(positions):
fw_positions[i] = simulation.four_warm_position(
self.sim_field_shape, self.pred_field_shape, pos)
sum_pos = np.sum(fw_positions[:, :, :], axis=0)
x[sample] = np.stack([field, sum_pos], axis=-1)
return np.array(x), mass
def build(self, samples=5000, mass=None):
"""Builds the dataset
:param samples: number of data points
:param mass: if set to None, N random masses are drawn between .1 and 1
:returns: x and y
"""
if mass is None:
mass = self.draw_mass(samples, self.rng)
x = []
ignored = 0
for sample in range(samples):
print(f"Produced {sample+1} samples\r", end="", flush=True)
try:
field, positions = self._simulate(self.rng, mass[sample])
if self.downsize:
field = fields.downsize_field(field, self.pred_field_shape)
pos = simulation.four_warm_position(
self.sim_field_shape, self.pred_field_shape,
positions[-1])
else:
pos = simulation.four_warm_position(
self.sim_field_shape, self.sim_field_shape,
positions[-1])
x.append(np.stack([field, pos], axis=-1))
except IndexError:
# Sometimes the particle achieves incredible acceleration and
# goes completely haywire
print('WARNING: 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), mass
def test_downsize_field():
field = np.array([[1, 1, 2, 2], [1, 1, 2, 2], [3, 3, 4, 4], [3, 3, 4, 4]])
expected = np.array([[1, 2], [3, 4]])
assert_array_equal(fields.downsize_field(field, (2, 2)), expected)