def _csr_swap_zero_nonzero_in_row(row: sparse.csr_matrix,
rng: np.random.Generator,
p: float = 0.1) -> sparse.csr_matrix:
"""
Swap 0 and nonzero values
Typically, most values are 0, so we do this for p * num_zero entries
This is exact, meaning we never swap fewer or more values
"""
assert row.shape[0] == 1
nonzero_idx = row.nonzero()[1]
arr = row.toarray().squeeze()
zero_idx = np.where(arr == 0)[0]
# Because # nonzero << # zero, we use # nonzero to determine number of swaps
n = int(round(len(nonzero_idx) * p))
# Choose indices to swap
zero_idx_swap = rng.choice(zero_idx, n, replace=False)
nonzero_idx_swap = rng.choice(nonzero_idx, n, replace=False)
# Transfer nonzero values to selected "aero" indices
arr[zero_idx_swap] = arr[nonzero_idx_swap]
# Zero out the original values at the nonzero indices
arr[nonzero_idx_swap] = 0
retval = sparse.csr_matrix(arr)
assert retval.shape == row.shape
assert len(retval.nonzero()[1]) == len(nonzero_idx)
return retval
def sim_markov_d(jump: la.lnarray,
peq: Optional[np.ndarray] = None,
num_jump: int = 10,
rng: np.random.Generator = RNG) -> la.lnarray:
"""Simulate Markov process trajectory.
Parameters
----------
jump : la.lnarray (n,n)
Discrete time stochastic matrix.
peq : la.lnarray (n,), optional
Initial-state distribution, default: use steady-state.
num_jump : int, optional, default: 10
Stop after this many jumps.
Returns
-------
states : la.lnarray (w,)
Vector of states visited.
"""
jump = la.asanyarray(jump)
if peq is None:
peq = calc_peq_d(jump)[0]
state_inds = la.arange(len(peq))
states_from = la.array(
[rng.choice(state_inds, size=num_jump - 1, p=p) for p in jump])
states = la.empty(num_jump)
states[0] = rng.choice(state_inds, p=peq)
for num in range(num_jump - 1):
states[num + 1] = states_from[states[num], num]
return states
def make_affiliation_network(f_coeffs: np.ndarray, g_coeffs: np.ndarray,
N: int, M: int,
rng: np.random.Generator) -> Network:
# get initial counts for how many stubs each agent and group has
agent_stubs: np.ndarray = rng.choice(len(f_coeffs), N, p=f_coeffs)
group_stubs: np.ndarray = rng.choice(len(g_coeffs), M, p=g_coeffs)
# balance the counts so that the groups can combine nicely
# This doesn't seem to be super necessary because it isn't feasible
# to balance all the time and the results look fine visually without it
# TODO: check that the actual results from the function match the results
# predicted by the math
timeout = 1_000
pre_calculed_indices = rng.choice(N + M, size=timeout)
while np.sum(agent_stubs) != np.sum(group_stubs) and timeout > 0:
stub_count_to_replace = pre_calculed_indices[timeout - 1]
if stub_count_to_replace < N:
agent_stubs[stub_count_to_replace] = rng.choice(len(f_coeffs),
1,
p=f_coeffs)
else:
group_stubs[stub_count_to_replace - N] = rng.choice(len(g_coeffs),
1,
p=g_coeffs)
timeout -= 1
group_membership = defaultdict(lambda: set())
# Assign agents to groups
# This assumes that agents can be assigned to the same group multiple times, but only
# one edge will appear
agents_with_stubs = np.nonzero(agent_stubs)[0]
groups_with_stubs = np.nonzero(group_stubs)[0]
while len(agents_with_stubs) > 0 and len(groups_with_stubs) > 0:
agent = rng.choice(agents_with_stubs)
group = rng.choice(groups_with_stubs)
group_membership[group].add(agent)
agent_stubs[agent] -= 1
group_stubs[group] -= 1
agents_with_stubs = np.nonzero(agent_stubs)[0]
groups_with_stubs = np.nonzero(group_stubs)[0]
# for agent_id, n_stubs in enumerate(agent_stubs):
# if len(np.nonzero(group_stubs)[0]) == 0:
# break
# for _ in range(n_stubs):
# possible_groups = np.nonzero(group_stubs)[0]
# if len(possible_groups) == 0:
# break
# group = rng.choice(possible_groups)
# group_membership[group].add(agent_id)
# group_stubs[group] -= 1
edges = it.chain.from_iterable(
it.combinations(agents, 2) for agents in group_membership.values())
G: nx.Graph = nx.empty_graph(N)
G.add_edges_from(edges)
return Network(G)
def _resample_ordinary_1(
sample: np.ndarray, size: int,
rng: np.random.Generator) -> _tp.Generator[np.ndarray, None, None]:
# i.i.d. sampling from empirical cumulative distribution of sample
n = len(sample)
for _ in range(size):
yield rng.choice(sample, size=n, replace=True)
def _ctmp_inverse(
n_samples: int, probs: np.ndarray, gamma: float,
csc_data: np.ndarray, csc_indices: np.ndarray,
csc_indptrs: np.ndarray,
rng: np.random.Generator) -> Tuple[Tuple[int], Tuple[int]]:
"""Apply CTMP algorithm to input counts dictionary, return
sampled counts and associated shots. Equivalent to Algorithm 1 in
Bravyi et al.
Args:
n_samples: Number of times to sample in CTMP algorithm.
probs: probability vector constructed from counts.
gamma: noise strength parameter
csc_data: Sparse CSC matrix data array (`csc_matrix.data`).
csc_indices: Sparse CSC matrix indices array (`csc_matrix.indices`).
csc_indptrs: Sparse CSC matrix indices array (`csc_matrix.indptrs`).
rng: RNG generator object.
Returns:
Tuple[Tuple[int], Tuple[int]]: Resulting list of shots and associated
signs from the CTMP algorithm.
"""
alphas = rng.poisson(lam=gamma, size=n_samples)
signs = np.mod(alphas, 2)
x_vals = rng.choice(len(probs), size=n_samples, p=probs)
# Apply CTMP sampling
r_vals = rng.random(size=alphas.sum())
y_vals = np.zeros(x_vals.size, dtype=np.int)
_markov_chain_compiled(y_vals, x_vals, r_vals, alphas, csc_data,
csc_indices, csc_indptrs)
return y_vals, signs
def make_array_with_random_subset_of_dna_seqs(
length: int,
num_seqs: int,
rng: np.random.Generator,
bases: Collection[str] = ('A', 'C', 'G', 'T')
) -> np.ndarray:
"""
Return 2D numpy array with random subset of size `num_seqs` of DNA sequences of given length.
Bases contains bases to be used: ('A','C','G','T') by default, but can be set to a subset of these.
Uses the encoding described in the documentation for DNASeqList. The result is a 2D array,
where each row represents a DNA sequence, and that row has one byte per base.
:param length: length of each row
:param num_seqs: number of rows
:param bases: DNA bases to use
:param rng: numpy random number generator (type returned by numpy.random.default_rng())
:return: 2D numpy array with random subset of size `num_seqs` of DNA sequences of given length
"""
if not set(bases) <= {'A', 'C', 'G', 'T'}:
raise ValueError(
f"bases must be a subset of {'A', 'C', 'G', 'T'}; cannot be {bases}"
)
if len(bases) == 0:
raise ValueError('bases cannot be empty')
elif len(bases) == 1:
raise ValueError('bases must have at least two elements')
base_bits = np.array([base2bits[base] for base in bases], dtype=np.ubyte)
arr = rng.choice(a=base_bits, size=(num_seqs, length))
unique_sorted_arr = np.unique(arr, axis=0)
return unique_sorted_arr
def sample_backwards(
fin_t: float,
fin_x: int,
gen: np.ndarray,
ome: np.random.Generator = np.random.default_rng()
) -> Skeleton:
"""
:param fin_t:
:param fin_x:
:param gen:
:param ome:
:return:
>>> # generate fixture
>>> dim = 3
>>> fin_t = 1e5
>>> gen = sample_trial_generator(dim)
>>> stat = get_stat(gen)
>>> fin_x = np.random.choice(dim, 1, p=get_stat(gen)).item()
>>> # test stationary distribution
>>> skel = sample_backwards(fin_t, fin_x, gen)
>>> probs = [np.sum(np.ediff1d(skel.t, to_end=(fin_t - skel.t[-1],))[skel.xt == i]) / fin_t for i in range(dim)]
>>> np.allclose(probs, stat, 1e-2)
True
"""
ftrans = get_trans(fin_t, gen)
btrans = ftrans[:, fin_x] / np.sum(ftrans[:, fin_x])
init_x = ome.choice(np.arange(len(btrans)), 1, p=btrans).item()
return sample_bridge(fin_t, init_x, fin_x, gen, ome)
def gravity_model_contact_events(agents: List[Agent],
positions: np.array,
env: simpy.Environment,
rng: np.random.Generator):
tree = KDTree(data=positions)
close_pairs = list(tree.query_pairs(r=contact_distance_upper_bound))
inverse_distances = np.array(
[np.linalg.norm(positions[idx1] - positions[idx2]) ** -contact_rate_gravity_exponent
for idx1, idx2 in close_pairs])
inverse_distances /= inverse_distances.sum()
while True:
choices = rng.choice(a=close_pairs, p=inverse_distances, size=len(agents)).tolist()
for choice in choices:
yield env.timeout(delay=rng.exponential(scale=1 / len(agents) / contact_rate_per_individual))
contact_agents = [agents[idx]
for idx in choice
if not agents[idx].state.symptomatic()
or rng.uniform() > p_symptomatic_individual_isolates]
if len(contact_agents) < 2:
# Symptomatic self-isolation means this is no longer a contact event and doesn't need
# recording. Skip to the next event.
continue
infected = get_infected(contact_agents, rng=rng)
for i in infected:
env.process(generator=infection_events(env=env, infected=i, rng=rng))
def try_make_new_random_edge(self, input_layer: Layer, output_layer: Layer,
rng: np.random.Generator) -> Optional[Edge]:
if type(output_layer) == OutputLayer:
return self.try_make_new_dense_edge(input_layer, rng)
mutation_operations = rng.choice(CnnGenome.EDGE_TYPE_FUNCTIONS,
len(CnnGenome.EDGE_TYPE_FUNCTIONS),
p=CnnGenome.EDGE_TYPE_PROBABILITIES,
replace=False)
for operation in mutation_operations:
edge = None
if operation == CnnGenome.add_conv_edge_mut:
edge = self.try_make_new_conv_edge(input_layer, output_layer,
rng)
elif operation == CnnGenome.add_separable_conv_edge_mut:
edge = self.try_make_new_separable_conv_edge(
input_layer, output_layer, rng)
elif operation == CnnGenome.add_pooling_edge_mut:
edge = self.try_make_new_pooling_edge(input_layer,
output_layer, rng)
if edge:
return edge
return None
def test_nonbonded_interaction_group_zero_interactions(
rng: np.random.Generator):
num_atoms = 33
num_atoms_ligand = 15
beta = 2.0
lamb = 0.1
cutoff = 1.1
box = 10.0 * np.eye(3)
conf = rng.uniform(0, 1, size=(num_atoms, 3))
ligand_idxs = rng.choice(num_atoms,
size=(num_atoms_ligand, ),
replace=False).astype(np.int32)
# shift ligand atoms in x by twice the cutoff
conf[ligand_idxs, 0] += 2 * cutoff
params = rng.uniform(0, 1, size=(num_atoms, 3))
potential = NonbondedInteractionGroup(
ligand_idxs,
np.zeros(num_atoms, dtype=np.int32),
np.zeros(num_atoms, dtype=np.int32),
beta,
cutoff,
)
du_dx, du_dp, du_dl, u = potential.unbound_impl(np.float64).execute(
conf, params, box, lamb)
assert (du_dx == 0).all()
assert (du_dp == 0).all()
assert du_dl == 0
assert u == 0
def make_env_attributes_task(
env: gym.Env,
task_params: Union[List[str], Dict[str, Any]],
seed: int = None,
rng: np.random.Generator = None,
noise_std: float = 0.2,
) -> Dict[str, Any]:
task: Dict[str, Any] = {}
rng: np.random.Generator = rng or np.random.default_rng(seed)
if isinstance(task_params, list):
task_params = {param: getattr(env.unwrapped, param) for param in task_params}
for attribute, default_value in task_params.items():
new_value = default_value
if isinstance(default_value, (int, float, np.ndarray)):
new_value *= rng.normal(1.0, noise_std)
# Clip the value to be in the [0.1*default, 10*default] range.
new_value = max(0.1 * default_value, new_value)
new_value = min(10 * default_value, new_value)
if isinstance(default_value, int):
new_value = round(new_value)
elif isinstance(default_value, bool):
new_value = rng.choice([True, False])
else:
raise NotImplementedError(
f"TODO: Don't yet know how to sample a random value for "
f"attribute {attribute} with default value {default_value} of type "
f" {type(default_value)}."
)
task[attribute] = new_value
return task
def sample_posterior(self,
nsamples: int,
nparams: int = 14,
rng: np.random.Generator = None) -> np.ndarray:
"""
Take a random sample from the posterior. This is useful for simulating
spectra for the purpose of illustrating how uncertainty in the model
parameters is reflected in the result.
Parameters
----------
nsamples : int
Number of random samples to draw from the posterior.
nparams : int, optional
Dimensionality of the model, i.e. number of parameters.
By default 14
rng : np.random.Generator, optional
Instance of a NumPy RNG, by default None, which
creates one.
Returns
-------
np.ndarray
Random samples drawn from the posterior.
"""
samples = (np.array(
self.posterior.posterior.to_array()).squeeze().reshape(
-1, nparams))
if rng is None:
rng = np.random.default_rng()
return rng.choice(samples, nsamples, axis=0)
def create_random_circuit(
n_qubits: int,
n_gates: int,
rng: np.random.Generator,
) -> Circuit:
"""Generates random circuit acting on nqubits with ngates for testing purposes.
The resulting circuit it saved to file in JSON format under 'circuit.json'.
Args:
n_qubits: The number of qubits in the circuit
n_gates: The number of gates in the circuit
rng: Numpy random generator
Returns:
Generated circuit.
"""
# Initialize all gates in set, not including RH or ZXZ
all_gates_lists = [
ONE_QUBIT_NO_PARAMS_GATES,
TWO_QUBITS_NO_PARAMS_GATES,
ONE_QUBIT_ONE_PARAM_GATES,
TWO_QUBITS_ONE_PARAM_GATES,
]
# Loop to add gates to circuit
circuit = Circuit()
for gate_i in range(n_gates):
# Pick gate type
gates_list = rng.choice(all_gates_lists)
gate = rng.choice(gates_list)
# Pick qubit to act on (control if two qubit gate)
qubits: Tuple[int, ...]
if gates_list in [ONE_QUBIT_NO_PARAMS_GATES, ONE_QUBIT_ONE_PARAM_GATES]:
index = rng.choice(range(n_qubits))
qubits = (int(index),)
else:
indices = rng.choice(range(n_qubits), size=2, replace=False)
qubits = tuple(int(i) for i in indices)
if gates_list in [ONE_QUBIT_ONE_PARAM_GATES, TWO_QUBITS_ONE_PARAM_GATES]:
param = rng.uniform(-np.pi, np.pi, size=1)
gate = gate(float(param))
circuit += gate(*qubits)
return circuit
def test_nonbonded_pair_list_interpolated_correctness(
ixn_group_size,
precision,
rtol,
atol,
cutoff,
beta,
lamb,
example_nonbonded_params,
example_conf,
example_box,
rng: np.random.Generator,
):
"Compares with jax reference implementation, with parameter interpolation."
num_atoms, _ = example_conf.shape
params = gen_params(example_nonbonded_params, rng)
# randomly select 2 interaction groups and construct all pairwise interactions
atom_idxs = rng.choice(
num_atoms,
size=(
2,
ixn_group_size,
),
replace=False,
).astype(np.int32)
pair_idxs = np.stack(np.meshgrid(atom_idxs[0, :],
atom_idxs[1, :])).reshape(2, -1).T
num_pairs, _ = pair_idxs.shape
scales = rng.uniform(0, 1, size=(num_pairs, 2))
lambda_plane_idxs = rng.integers(-2, 3, size=(num_atoms, ), dtype=np.int32)
lambda_offset_idxs = rng.integers(-2,
3,
size=(num_atoms, ),
dtype=np.int32)
ref_potential = nonbonded.interpolated(
make_ref_potential(pair_idxs, scales, lambda_plane_idxs,
lambda_offset_idxs, beta, cutoff))
test_potential = NonbondedPairListInterpolated(pair_idxs, scales,
lambda_plane_idxs,
lambda_offset_idxs, beta,
cutoff)
GradientTest().compare_forces(
example_conf,
params,
example_box,
lamb,
ref_potential,
test_potential,
precision=precision,
rtol=rtol,
atol=atol,
)
def _resample_ordinary_n(
samples: _tp.List[np.ndarray], size: int,
rng: np.random.Generator) -> _tp.Generator[np.ndarray, None, None]:
n = len(samples[0])
indices = np.arange(n)
for _ in range(size):
m = rng.choice(indices, size=n, replace=True)
yield tuple(s[m] for s in samples)
def generate_random_marker(
depth: int,
max_branches: int,
rng: np.random.Generator,
constant_condition_text: Optional[Text],
possible_conditions: List[Type[ConditionMarker]],
constant_negated: Optional[bool],
possible_operators: List[Type[OperatorMarker]],
) -> Tuple[Marker, int]:
"""Generates an (max_branches)-ary tree with the specified depth."""
if depth == 0:
condition_class = possible_conditions[rng.choice(
len(possible_conditions))]
negated = bool(
rng.choice(2)) if constant_negated is None else constant_negated
condition_text = constant_condition_text or f"{rng.choice(1000)}"
return condition_class(text=condition_text, negated=negated), 1
else:
negated = bool(
rng.choice(2)) if constant_negated is None else constant_negated
operator_class = possible_operators[rng.choice(
len(possible_operators))]
num_branches = operator_class.expected_number_of_sub_markers()
if num_branches is None:
num_branches = rng.choice(max_branches - 1) + 1
marker_size = 0
sub_markers = []
for _ in range(num_branches):
sub_marker, sub_marker_size = generate_random_marker(
depth=depth - 1,
max_branches=max_branches,
rng=rng,
constant_negated=constant_negated,
constant_condition_text=constant_condition_text,
possible_operators=possible_operators,
possible_conditions=possible_conditions,
)
marker_size += sub_marker_size
sub_markers.append(sub_marker)
marker = operator_class(markers=sub_markers, negated=negated)
marker_size += 1
return marker, marker_size
def _csr_swap_in_row(row: sparse.csr_matrix,
rng: np.random.Generator,
p: float = 0.1) -> sparse.csr_matrix:
"""
Helper function for swapping nonzero values in a given row
"""
assert row.shape[0] == 1, f"Did not get a row!"
nonzero_idx = row.nonzero()[1]
shuffle_idx = np.arange(len(nonzero_idx))
# Randomly choose a proportion of the nonzero indices to shuffle
n = int(round(len(shuffle_idx) * p))
swap_idx = nonzero_idx[rng.choice(shuffle_idx, size=n, replace=False)]
# Shuffle the indices we chose above
dest_idx = rng.choice(swap_idx, size=len(swap_idx), replace=False)
assert swap_idx.shape == dest_idx.shape
arr = row.toarray().squeeze()
assert np.all(arr[swap_idx] != 0)
arr[dest_idx] = arr[swap_idx]
retval = sparse.csr_matrix(arr)
return retval
def _generate_random_example_for_one_session_and_one_marker(
rng: np.random.Generator, ) -> Tuple[List[EventMetaData], List[int]]:
"""Generates a random marker extraction result for a single session and marker.
Args:
rng: a random number generator
Returns:
the event list representing the marker extraction result as well as
the plain list of numbers used as "preceding user turns" in that extraction
result
"""
applied = int(rng.choice(10))
all_preceding_user_turn_numbers = [
int(rng.choice(20)) for _ in range(applied)
]
event_list = [
EventMetaData(idx=int(rng.choice(100)),
preceding_user_turns=preceding_user_turns)
for preceding_user_turns in all_preceding_user_turn_numbers
]
return event_list, all_preceding_user_turn_numbers
def _iter_random_indices(
rng: np.random.Generator,
num_sources: int,
random_batch_size=1000,
p: Optional[List[float]] = None,
) -> Iterator[int]:
"""Get an infinite iterator that randomly samples the index of the source to pick examples from."""
if p is None:
while True:
yield from (int(i) for i in rng.integers(0, num_sources, size=random_batch_size))
else:
while True:
yield from (int(i) for i in rng.choice(num_sources, size=random_batch_size, p=p))
def random_subspace(all_features: list, k: int, rng: np.random.Generator):
"""Utility function to generate a random feature subspace of length k
Parameters
----------
all_features
List of possible features to select from.
k
Subspace length.
rng
Random number generator (initialized).
"""
return rng.choice(all_features, k, replace=False)
def select(
population: np.ndarray,
scores: np.ndarray,
indexes: np.ndarray,
rand_generator: np.random.Generator = _DEFAULT_RAND_GEN) -> np.ndarray:
"""
Inputs:
population
-- array of chromosomes
-- each chromosome must be represented as 1D numpy boolean array
scores
-- array of fitness values corresponding to chromosomes
-- each fitness value must be a non-negative real number
indexes
-- 1D numpy integer array
-- indexes of chromosomes in the population (row index)
rand_generator
-- instance of Numpy Random Generator
-- Generator with Numpy default BitGenerator (PCG64)
and None as a seed value is used as default
Selection is based on the roulette wheel method (fitness proportionate selection)
where probability of a chromosome being selected is related to its fitness value.
This does not guarantee that the best chromosome sequence (binary pattern)
will be selected but helps to avoid local optima.
Returns nested (2D) numpy boolean array, the entire next generation of chromosomes
(solution candidates) chosen with repetition from a given population.
"""
probabilities = scores / np.sum(scores)
indexes = rand_generator.choice(indexes,
size=indexes.size,
replace=True,
p=probabilities)
rand_generator.shuffle(indexes)
return population[indexes]
def add_edge_mut(self, rng: np.random.Generator) -> bool:
"""
This performs an add edge mutation by randomly selecting two layers and trying to create an edge
between them. If an edge cannot be created, two different layers will be selected.
This process will be repeated until an edge is successfully created.
"""
logging.info("attempting add_edge mutation")
mutation_operations = rng.choice(CnnGenome.EDGE_TYPE_FUNCTIONS,
len(CnnGenome.EDGE_TYPE_FUNCTIONS),
p=CnnGenome.EDGE_TYPE_PROBABILITIES,
replace=False)
for operation in mutation_operations:
if operation(self, rng):
logging.info("successfully completed add_edge mutation")
return True
logging.info("failed to complete add_edge mutation")
return False
def generate_random_from_kernel(a: np.ndarray, d: int,
rng: np.random.Generator) -> np.ndarray:
""" Use the random number generator `rng` to draw a random `d`-dimensional
subspace from the kernel of matrix `a`. If `dim(ker(a)) < d`, an exception
is raised.
"""
evals, evecs = np.linalg.eigh(a.T @ a)
# kernel is where evals are (almost) zero
ker_dims = (evals < 1e-10).nonzero()[0]
size_ker = len(ker_dims)
if size_ker < d:
raise ValueError(
"Kernel dimension lower than requested subspace dimension")
chosen_dims = rng.choice(ker_dims, size=d, replace=False)
chosen_basis = evecs[:, chosen_dims]
b = rng.normal(size=(d, d)) @ chosen_basis.T
return b
def resample_and_perturbate(particles: List[Particle],
weights: List[float],
rng: np.random.Generator) -> Particle:
""" Resampling part of ABC-SMC, selects randomly a new particle from the
list of previously accepted. Then perturb it slightly. This causes
the persistence and selection of fit particles in a manner similar to
evolution algorithms. The perturbation is done until the candidate is valid,
this is because the perturbation is unaware of the support of the distribution
of the particle.
:param particles: List of particles from previous population
:param weights: List of weights of particles from previous population
:param rng: Random state for random sampling into particles
:return: Newly created particles sampled from previous population and perturbated
"""
while True:
particle = rng.choice(particles, p=weights / np.sum(weights))
particle = particle.generate_perturbated()
if particle.validate_particle():
return particle
def randomlyInfectRegions(
network: NetworkOfPopulation, regions: int, age_groups: List[Age],
infected: float,
random_state: np.random.Generator) -> Dict[NodeName, Dict[Age, float]]:
"""Randomly infect regions to initialize the random simulation
:param network: object representing the network of populations
:param regions: The number of regions to expose.
:param age_groups: Age groups to infect
:param infected: People to infect
:param random_state: Random state for random number generation
:return: Structure of initially infected regions with number
"""
infections: Dict[NodeName, Dict[Age, float]] = {}
for regionID in random_state.choice(list(network.graph.nodes()),
size=regions):
infections[regionID] = {}
for age in age_groups:
infections[regionID][age] = infected
return infections
def simulate(rng: np.random.Generator,
policy: callable,
T: int,
s0: int = 5,
D: np.ndarray = D,
P: np.ndarray = P):
"""Simulate widget sales for a given policy."""
s = np.zeros(T + 1) # states
a = np.zeros(T) # actions
r = np.zeros(T) # rewards
s[0] = s0 # initial state
for t in tqdm(range(T)):
# Sample demand
d = rng.choice(D, p=P)
# Record action, reward, and next state
a[t] = policy(s[t])
r[t] = reward(s[t], a[t], d)
s[t + 1] = transition(s[t], a[t], d)
return s, a, r
def sample_forwards(
fin_t: float,
init_x: Optional[int],
gen: np.ndarray,
ome: np.random.Generator = np.random.default_rng()
) -> Skeleton:
"""
:param fin_t:
:param init_x:
:param gen:
:param ome:
:return:
>>> # generate fixture
>>> dim = 3
>>> fin_t = 1e5
>>> gen = sample_trial_generator(dim)
>>> stat = get_stat(gen)
>>> # test stationary distribution
>>> skel = sample_forwards(fin_t, None, gen)
>>> probs = [np.sum(np.ediff1d(skel.t, to_end=(fin_t - skel.t[-1],))[skel.xt == i]) / fin_t for i in range(dim)]
>>> np.allclose(probs, stat, 1e-2)
True
"""
dim = gen.shape[0]
if init_x is None:
init_x = ome.choice(dim, 1, p=get_stat(gen)).item()
t, x = [0], [init_x]
while t[-1] < fin_t:
hold = ome.exponential(1 / np.delete(gen[x[-1]], x[-1]))
move = np.argmin(hold)
t.append(t[-1] + hold[move])
x.append(move + int(move >= x[-1]))
return Skeleton(np.array(t[:-1]), np.array(x[:-1]), fin_t)
def pick_randomly(
universe: pd.DataFrame,
holdings_number: pd.Series,
rng: np.random.Generator,
) -> pd.DataFrame:
universe, holdings_number = align(universe, holdings_number)
universe_array = universe.to_numpy()
holdings_number_array = holdings_number.to_numpy()
choices = np.zeros_like(universe_array, dtype=bool)
for i in range(0, len(choices)):
tradable = np.where(universe_array[i])[0]
choice = rng.choice(
tradable,
size=holdings_number_array[i],
replace=False,
)
choices[i, choice] = True
return pd.DataFrame(
choices,
index=universe.index.copy(),
columns=universe.columns.copy(),
)
def test_nonbonded_interaction_group_consistency_allpairs_constant_shift(
num_atoms,
num_atoms_ligand,
precision,
rtol,
atol,
cutoff,
beta,
lamb,
example_nonbonded_params,
example_conf,
example_box,
rng: np.random.Generator,
):
"""Compares with reference nonbonded_v3 potential, which computes
the sum of all pairwise interactions. This uses the identity
U(x') - U(x) = U_AB(x') - U_AB(x)
where
- U is the all-pairs potential over all atoms
- U_A, U_B are all-pairs potentials for interacting groups A and
B, respectively
- U_AB is the "interaction group" potential, i.e. the sum of
pairwise interactions (a, b) where "a" is in A and "b" is in B
- the transformation x -> x' does not affect U_A or U_B (e.g. a
constant translation applied to each atom in one group)
"""
conf = example_conf[:num_atoms]
params = example_nonbonded_params[:num_atoms, :]
lambda_plane_idxs = rng.integers(-2, 3, size=(num_atoms, ), dtype=np.int32)
lambda_offset_idxs = rng.integers(-2,
3,
size=(num_atoms, ),
dtype=np.int32)
def ref_allpairs(conf):
return prepare_reference_nonbonded(
params=params,
exclusion_idxs=np.array([], dtype=np.int32),
scales=np.zeros((0, 2), dtype=np.float64),
lambda_plane_idxs=lambda_plane_idxs,
lambda_offset_idxs=lambda_offset_idxs,
beta=beta,
cutoff=cutoff,
)(conf, params, example_box, lamb)
ligand_idxs = rng.choice(num_atoms,
size=(num_atoms_ligand, ),
replace=False).astype(np.int32)
def test_ixngroups(conf):
_, _, _, u = (NonbondedInteractionGroup(
ligand_idxs,
lambda_plane_idxs,
lambda_offset_idxs,
beta,
cutoff,
).unbound_impl(precision).execute(conf, params, example_box, lamb))
return u
conf_prime = np.array(conf)
conf_prime[ligand_idxs] += rng.normal(0, 0.01, size=(3, ))
ref_delta = ref_allpairs(conf_prime) - ref_allpairs(conf)
test_delta = test_ixngroups(conf_prime) - test_ixngroups(conf)
np.testing.assert_allclose(ref_delta, test_delta, rtol=rtol, atol=atol)
def test_nonbonded_interaction_group_consistency_allpairs_lambda_planes(
num_atoms,
num_atoms_ligand,
precision,
rtol,
atol,
cutoff,
beta,
lamb,
example_nonbonded_params,
example_conf,
example_box,
rng: np.random.Generator,
):
"""Compares with reference nonbonded_v3 potential, which computes
the sum of all pairwise interactions. This uses the identity
U = U_A + U_B + U_AB
where
- U is the all-pairs potential over all atoms
- U_A, U_B are all-pairs potentials for interacting groups A and
B, respectively
- U_AB is the "interaction group" potential, i.e. the sum of
pairwise interactions (a, b) where "a" is in A and "b" is in B
U is computed using the reference potential over all atoms, and
U_A + U_B computed using the reference potential over all atoms
separated into 2 lambda planes according to which interacting
group they belong
"""
conf = example_conf[:num_atoms]
params = example_nonbonded_params[:num_atoms, :]
max_abs_offset_idx = 2 # i.e., lambda_offset_idxs in {-2, -1, 0, 1, 2}
lambda_offset_idxs = rng.integers(-max_abs_offset_idx,
max_abs_offset_idx + 1,
size=(num_atoms, ),
dtype=np.int32)
def make_reference_nonbonded(lambda_plane_idxs):
return prepare_reference_nonbonded(
params=params,
exclusion_idxs=np.array([], dtype=np.int32),
scales=np.zeros((0, 2), dtype=np.float64),
lambda_plane_idxs=lambda_plane_idxs,
lambda_offset_idxs=lambda_offset_idxs,
beta=beta,
cutoff=cutoff,
)
ref_allpairs = make_reference_nonbonded(np.zeros(num_atoms,
dtype=np.int32))
ligand_idxs = rng.choice(num_atoms,
size=(num_atoms_ligand, ),
replace=False).astype(np.int32)
# for reference U_A + U_B computation, ensure minimum distance
# between a host and ligand atom is at least one cutoff distance
# when lambda = 1
lambda_plane_idxs = np.zeros(num_atoms, dtype=np.int32)
lambda_plane_idxs[ligand_idxs] = 2 * max_abs_offset_idx + 1
ref_allpairs_minus_ixngroups = make_reference_nonbonded(lambda_plane_idxs)
def ref_ixngroups(*args):
return ref_allpairs(*args) - ref_allpairs_minus_ixngroups(*args)
test_ixngroups = NonbondedInteractionGroup(
ligand_idxs,
np.zeros(num_atoms, dtype=np.int32), # lambda plane indices
lambda_offset_idxs,
beta,
cutoff,
)
GradientTest().compare_forces(
conf,
params,
example_box,
lamb=lamb,
ref_potential=ref_ixngroups,
test_potential=test_ixngroups,
rtol=rtol,
atol=atol,
precision=precision,
)