3.3. wavesim.multidomain

class Domain(pixel_size, shape, device)

Bases: object

Base class for all simulation domains

This base class defines the interface for operations that are common for all simulation types, and for MultiDomain. todo: the design using slots minimizes memory use, but it is a suboptimal design because it mixes mutable and immutable state. This design should be revisited so that the Domain is immutable, and the code that runs the algorithms performs the memory management.

abstractmethod add_source(slot, weight)

abstractmethod clear(slot)

Clears the data in the specified slot

coordinates(dim, type='linear')

Returns the real-space coordinates along the specified dimension, starting at 0

coordinates_f(dim)

Returns the Fourier-space coordinates along the specified dimension

abstractmethod create_empty_vdot()

Create an empty tensor for the Vdot tensor

abstractmethod get(slot, copy=False)

Returns the data in the specified slot.

Parameters:

• slot (int) – slot from which to return the data

• copy – if True, returns a copy of the data. Otherwise, may return the original data possible.

Note that this data may be overwritten by the next call to domain.

abstractmethod inner_product(slot_a, slot_b)

Computes the inner product of two data vectors

Note

The vectors may be represented as multidimensional arrays, but these arrays must be contiguous for this operation to work. Although it would be possible to use flatten(), this would create a copy when the array is not contiguous, causing a hidden performance hit.

abstractmethod inverse_propagator(slot_in, slot_out)

Applies the operator (L+1) x .

This operation is not needed for the Wavesim algorithm, but is provided for testing purposes, and can be used to evaluate the residue of the solution.

abstractmethod medium(slot_in, slot_out, mnum)

Applies the operator 1-Vscat.

abstractmethod mix(weight_a, slot_a, weight_b, slot_b, slot_out)

Mixes two data arrays and stores the result in the specified slot

abstractmethod propagator(slot_in, slot_out)

Applies the operator (L+1)^-1 x.

abstractmethod set(slot, data)

Copy the date into the specified slot

abstractmethod set_source(source)

Sets the source term for this domain.

class HelmholtzDomain(permittivity, periodic, pixel_size=0.25, wavelength=None, n_boundary=0, n_slots=2, stand_alone=True, Vwrap=None, device=None, debug=False)

Bases: Domain

Represents a single domain of the simulation.

The Domain object encapsulates all data that is stored on a single computation node (e.g. a GPU or a node in a cluster), and provides methods to perform the basic operations that the Wavesim algorithm needs.

Note

Domain currently works only for the Helmholtz equation and the PyTorch backend. If we want to have further functionality, this class should be refactored.

add_source(slot, weight)

Adds the source term to the data in the specified slot.

apply_corrections(wrap_corrections, transfer_corrections, slot)

Apply wrapping/transfer corrections

Transfer corrections correspond to a contribution from neighboring domains, and are added to the current domain. Wrap corrections correct for the periodicity of the fft. They are subtracted from the domain. In this case, there is an additional factor of -1 because this function is called from medium, which applies 1-V instead of V. Therefore, transfer corrections are now subtracted, and wrap corrections are added.

Parameters:

• slot (int) – slot index for the data to which the corrections are applied. Operation is always in-place

• wrap_corrections – list of 6 corrections for wrap-around (may contain None for periodic boundary)

• transfer_corrections – list of 6 corrections coming from neighboring segments (may contain None for end of domain)

clear(slot)

Clears the data in the specified slot

compute_corrections(slot_in)

Computes the edge corrections by multiplying the first and last pixels of each line with the Vwrap matrix.

The corrections are stored in self.edges. TODO: re-use this memory

create_empty_vdot()

Creates an empty tensor to force the allocation of memory for the Vdot tensor.

This tensor is always 8.1 MiB, irrespective of the domain size.

get(slot, copy=False)

Returns the data in the specified slot.

Parameters:

• slot (int) – slot from which to return the data

• copy – if True, returns a copy of the data. Otherwise, may return the original data possible. Note that this data may be overwritten by the next call to domain.

initialize_scale(scale)

Scales all operators.

Computes Bscat (from the temporary storage 0), the propagator kernel (from the temporary value in propagator_kernel), and scales Vwrap.

scale

Scaling factor of the problem. Its magnitude is chosen such that the operator V = scale · (the scattering potential + the wrapping correction) has norm < 1. The complex argument is chosen such that L+V is accretive.

initialize_shift(shift)

Shifts the scattering potential and propagator kernel, then returns the norm of the shifted operator.

inner_product(slot_a, slot_b)

Computes the inner product of two data vectors

Note

The vectors may be represented as multidimensional arrays, but these arrays must be contiguous for this operation to work. Although it would be possible to use flatten(), this would create a copy when the array is not contiguous, causing a hidden performance hit.

inverse_propagator(slot_in, slot_out)

Applies the operator (L+1) x .

This operation is not needed for the Wavesim algorithm, but is provided for testing purposes, and can be used to evaluate the residue of the solution.

medium(slot_in, slot_out, mnum=None)

Applies the operator 1-Vscat.

Note

Does not apply the wrapping correction. When part of a multi-domain, the wrapping correction is applied by the medium() function of the multi-domain object and this function should not be called directly.

mix(weight_a, slot_a, weight_b, slot_b, slot_out)

Mixes two data arrays and stores the result in the specified slot

propagator(slot_in, slot_out)

Applies the operator (L+1)^-1 x.

set(slot, data)

Copy the date into the specified slot

set_source(source)

Sets the source term for this domain.

class MultiDomain(permittivity, periodic, pixel_size=0.25, wavelength=None, n_domains=(1, 1, 1), n_boundary=8, device=None, debug=False)

Bases: Domain

Class for generating medium (B) and propagator (L+1)^(-1) operators, scaling, and setting up wrapping and transfer corrections

add_source(slot, weight)

Add the source to the field in slot

clear(slot)

Clear the field in the specified slot

create_empty_vdot()

Create an empty tensor for the Vdot tensor

get(slot, copy=False, device=None)

Get the field in the specified slot, this gathers the fields from all subdomains and puts them in one big array

Parameters:

• slot (int) – slot to get the data from

• copy – if True, returns a copy of the data. Otherwise, may return the original data possible. Note that this data may be overwritten by the next call to domain.

• device – device on which to store the data. Defaults to the primary device

inner_product(slot_a, slot_b)

Compute the inner product of the fields in slots a and b

Note

Use sqrt(inner_product(slot_a, slot_a)) to compute the norm of the field in slot_a. There is a large but inconsistent difference in performance between vdot and linalg.norm. Execution time can vary a factor of 3 or more between the two, depending on the input size and whether the function is executed on the CPU or the GPU.

inverse_propagator(slot_in, slot_out)

Apply inverse propagator operators L+1 to subdomains/patches of x.

medium(slot_in, slot_out, mnum=None)

Apply the medium operator B, including wrapping corrections.

Parameters:

• slot_in (int) – slot holding the input field

• slot_out (int) – slot that will receive the result

• mnum – # of the medium() call in preconditioned iteration. 0 for first, 1 for second medium call.

mix(weight_a, slot_a, weight_b, slot_b, slot_out)

Mix the fields in slots a and b and store the result in slot_out

propagator(slot_in, slot_out)

Apply propagator operators (L+1)^-1 to subdomains/patches of x.

set(slot, data)

Copy the date into the specified slot

set_source(source)

Split the source into subdomains and store in the subdomain states.

combine(domains, device='cpu')

Concatenates a 3-d array of 3-d tensors

empty_cache()

Release all unoccupied cached memory currently held by the caching allocator so that those can be used in other GPU application and visible in nvidia-smi.

Note

empty_cache() doesn’t increase the amount of GPU memory available for PyTorch. However, it may help reduce fragmentation of GPU memory in certain cases. See cuda-memory-management for more details about GPU memory management.

is_zero(x)

Check if x is zero

Some functions allow specifying 0 or 0.0 instead of a torch tensor, to indicate that the array should be cleared. This function returns True if x is a scalar 0 or 0.0. It raises an error if x is a scalar that is not equal to 0 or 0.0, and returns False otherwise.

list_to_array(input, depth)

Convert a nested list of depth depth to a numpy object array

partition(array, n_domains)

Split a 3-D array into a 3-D set of sub-arrays of approximately equal sizes.