sample_builder.sample_generator

User level functions for generating samples.

FIXME: write proper doc-string

@private

Copyright © 2022-2025 HQS Quantum Simulations GmbH. All Rights Reserved.

def generate_mask(positions: numpy.ndarray, attrs: dict) -> numpy.ndarray:

Generate a mask.

Arguments:
  • positions (np.ndarray): Sites' positions.
  • attrs (dict): dictionary specifying the mask to generate.
Raises:
  • ValueError: if attrs["shape"] is not valid.
  • ValueError: if the shape of the positions is not valid.
Returns:

np.ndarray: array of the same shape of r, selecting the wanted sites.

def generate_modulation_profile( positions: numpy.ndarray, direction: Union[list[float], numpy.ndarray], sampled_distributions: Union[list, numpy.ndarray], x_distribution: Union[list, numpy.ndarray, NoneType] = None, order: int = 1) -> numpy.ndarray:

Generate a modulation from a provided sampled distribution profile.

Arguments:
  • positions (np.ndarray): Sites positions.
  • direction (Union[list[float], np.ndarray]): Direction in which to generate the stacks.
  • sampled_distributions (Union[list, np.ndarray]): sampled distribution profile for each species. Columns should contain the sampled distribution for each species.
  • x_distribution: (Optional[Union[list, np.ndarray]]): The x-coordinates of the data points of the sampled distributions. The same x-coordinates will be used for all distributions.
  • order: (int): The order of the interpolation. Defaults to 1 (linear interpolation).
Returns:

np.ndarray: modulation according to the provided sampled profiles.

def generate_modulation_uniform( positions: numpy.ndarray, distribution: Union[list, numpy.ndarray]) -> numpy.ndarray:

Generate a uniform modulation.

Arguments:
  • positions (np.ndarray): Sites positions.
  • distribution (Union[list, np.ndarray]): Distribution to be assigned to each site.
Returns:

np.ndarray: Uniform modulation.

def generate_modulation_stacks( positions: numpy.ndarray, direction: Union[list, numpy.ndarray], widths: Union[list, numpy.ndarray], distributions: Union[list, numpy.ndarray], transition_width: float = 0) -> numpy.ndarray:

Generate a stacks modulation.

The system will be divided in different stacks along the selected direction and each stack will have the probability distribution defined in the homonymous argument.

Arguments:
  • positions (np.ndarray): Sites' positions.
  • direction (Union[list, np.ndarray]): Direction in which to generate the stacks.
  • widths (Union[list, np.ndarray]): Widths of the different stacks.
  • distributions (Union[list, np.ndarray]): Distributions for each stack in a nested list, e.g., [[0.5, 0.5], [1, 0]] for two stacks and two species.
  • transition_width (float): Width of the smooth transition between stacks. Defaults to 0.
Returns:

np.ndarray: Stacks modulation.

def generate_modulation_linear( positions: numpy.ndarray, direction: Union[list, numpy.ndarray], selected_species: int, n_species: int, start_value: float, stop_value: float, from_unit_cell_idx: Union[list, numpy.ndarray, NoneType] = None, to_unit_cell_idx: Union[list, numpy.ndarray, NoneType] = None) -> numpy.ndarray:

Generate linear modulation.

The probability of the selected species will be modulated linearly along the specified direction according to the arguments passed. The rest of the probability (up to a total of 1) will be distributed uniformly among the other species.

Arguments:
  • positions (np.ndarray): Sites' positions.
  • direction (Union[list, np.ndarray]): Direction of the linear modulation.
  • selected_species (int): Species which should be linearly modulated.
  • n_species (int): Total number of species to be generated.
  • start_value (float): Starting modulation value, should be in [0, 1].
  • stop_value (float): Final modulation value, should be in [0, 1].
  • from_unit_cell_idx (Optional[Union[list, np.ndarray]]): Unit cell index from which to start the linear modulation. Should have format [x_idx, y_idx, z_idx].
  • to_unit_cell_idx (Optional[Union[list, np.ndarray]]): Last unit cell index of the linear modulation. Should have format [x_idx, y_idx, z_idx].
Raises:
  • ValueError: If startValue or stopValue are not within [0, 1].
Returns:

np.ndarray: linear modulation.

def generate_modulation_sin( positions: numpy.ndarray, direction: Union[list, numpy.ndarray], selected_species: Union[list, int], n_species: int, y0: Union[list, float], amplitude: Union[list, float], wavelength: Union[list, float], phi: Union[list, int, NoneType] = None) -> numpy.ndarray:

Generate sinusoidal modulation of one or more species.

The probability of the selected species (single or multiple) will be modulated as

.. math: y0 + amplitude sin(\frac{2 \pi}{wavelength} (pp - phi))

where pp are the positions projected in the selected direction. The rest of the probability (up to a total of 1) will be distributed uniformly among the rest of the species.

Arguments:
  • positions (np.ndarray): Sites' positions.
  • direction (Union[list, np.ndarray]): Direction of the sinusoidal modulation.
  • selected_species (Union[list, int]): Single or multiple species to modulate.
  • n_species (int): Total number of species to be generated.
  • y0 (Union[list, float]): Average value(s) of the modulation(s).
  • amplitude (Union[list, float]): Amplitude(s) of the modulation(s).
  • wavelength (Union[list, float]): Wavelength(s) of the modulation(s).
  • phi (Optional[Union[list, int]]): Phase shift(s) of the modulation(s).
Raises:
  • ValueError: if any inconsistency with the input is detected.
Returns:

np.ndarray: Sinusoidal modulation.

def generate_modulation_cos( positions: numpy.ndarray, direction: Union[list, numpy.ndarray], selected_species: Union[list, int], n_species: int, y0: Union[list, float], amplitude: Union[list, float], wavelength: Union[list, float], phi: Union[list[float], float, NoneType] = None) -> numpy.ndarray:

Generate cosinusoidal modulation of one or more species.

The probability of the selected species (single or multiple) will be modulated as

.. math: y0 + amplitude cos(\frac{2 \pi}{wavelength} (pp - phi))

where pp are the positions projected in the selected direction. The rest of the probability (up to a total of 1) will be distributed uniformly among the rest of the species.

Arguments:
  • positions (np.ndarray): Sites' positions.
  • direction (Union[list, np.ndarray]): Direction of the sinusoidal modulation.
  • selected_species (Union[list, int]): Single or multiple species to modulate.
  • n_species (int): Total number of species to be generated.
  • y0 (Union[list, float]): Average value(s) of the modulation(s).
  • amplitude (Union[list, float]): Amplitude(s) of the modulation(s).
  • wavelength (Union[list, float]): Wavelength(s) of the modulation(s).
  • phi (Optional[Union[list, int]]): Phase shift(s) of the modulation(s).
Raises:
  • ValueError: if any inconsistency with the input is detected.
Returns:

np.ndarray: Cosinusoidal modulation.

def generate_choice(site_dist: numpy.ndarray, rng: numpy.random._generator.Generator) -> int:

Generate random choice given the probability distribution.

Generates an integer between 0 and N, with the given probabilities. N is the number of probabilities passed.

Arguments:
  • site_dist (np.ndarray): probability distribution for the choice
  • rng (np.random.Generator): random number generator to be used.
Returns:

int: random choice given the probabilities.

def generate_sample( site_dist: numpy.ndarray, rng: numpy.random._generator.Generator) -> numpy.ndarray:

Generate random sample given the probability distribution.

For each site, generates an integer between 0 and N, with the given probabilities. N is the number of probabilities for each site.

Arguments:
  • site_dist (np.ndarray): probability distribution for the whole sample.
  • rng (np.random.Generator): random number generator to be used.
Returns:

np.ndarray: random choice given the probabilities.

def combine_mask_modulation(mods: list, masks: list) -> numpy.ndarray:

Combine masks and modulations while performing consistency checks.

Arguments:
  • mods (list): list of modulations to combine.
  • masks (list): list of corresponing masks.
Raises:
  • ValueError: if the masks are not complementary.
Returns:

np.ndarray: combined modulation.

def combine_modulations( mods: list, operation: str, default: Optional[list] = None) -> numpy.ndarray:

Combine modulations according to the specified operation.

Arguments:
  • mods (list): list of modulations to combine.
  • operation (str): type of operation used for mixing.
  • default (Optional[list]): default distribution for cases where NaNs appear.
Raises:
  • ValueError: if the operation is not valid.
Returns:

np.ndarray: combined modulation.

def sample_view( sample: numpy.ndarray, window_pos: Union[list, numpy.ndarray], window_size: Union[list, numpy.ndarray]) -> numpy.ndarray:

Generate a view of the provided sample.

NOTE: not the most general. Silently assumes sample is of the shape (x, y, z, unit_cell).

Arguments:
  • sample (np.ndarray): Sample.
  • window_pos (Union[list, np.ndarray]): Position (x, y, z) of the selecting window.
  • window_size (Union[list, np.ndarray]): Size of the window.
Raises:
  • ValueError: if the provided window parameters result in an out-of-bounds window.
Returns:

np.ndarray: View over the original sample.