Learned Representations

The pythia.learned module implements building blocks for creating learned representations of particle systems. They are implemented within the keras framework for ease of integration into standard workflows.

class pythia.learned.bonds.InertiaRotation(num_rotations=1, initial_mass_variance=0.25, center=False, **kwargs)

Generate rotation-invariant point clouds by orienting via principal axes of inertia

InertiaRotation takes an array of neighborhood points (shape (…, num_neighbors, 3)) and outputs one or more copies which have been rotated according to the principal axes of inertia of the neighborhood. It does this using masses that can be varied for each point and each rotation.

For an input of shape (…, num_neighbors, 3), InertiaRotation produces an output of shape (…, num_rotations, num_neighbors, 3).

Before computing the inertia tensor, points can optionally be centered via the center argument. A value of True centers the points as if all masses were 1, a value of “com” centers the points using the learned masses, and a value of False (the default) does not center at all.

Parameters:

• num_rotations – number of rotations to produce
• initial_mass_variance – Variance of the initial mass distribution (mean 1)
• center – Center the mass points before computing the inertia tensor (see description above)

class pythia.learned.spherical_harmonics.SphericalHarmonics(lmax, negative_m=False, **kwargs)

Compute the (complex) spherical harmonic decomposition given a set of cartesian coordinates

For an input of shape (…, 3), SphericalHarmonics will produce an output of shape (…, num_sphs), where num_sphs is the number of spherical harmonics produced given the lmax and negative_m arguments.

Parameters:

• lmax – maximum spherical harmonic degree to compute
• negative_m – If True, consider m=-l to m=l for each spherical harmonic l; otherwise, consider m=0 to m=l

class pythia.learned.spherical_harmonics.NeighborAverage(*args, **kwargs)

Compute a weighted average an array of complex-valued spherical harmonics over all points in a neighborhood

Given an input of shape (…, num_rotations, num_neighbors, num_sphs), NeighborAverage will produce an output of shape (…, num_rotations, num_sphs). Each neighbor from each rotation is assigned a learnable weight to contribute to the sum.

class pythia.learned.spherical_harmonics.ComplexProjection(num_projections=1, conversion='abs', activation=None, kernel_initializer='glorot_uniform', bias_initializer='random_normal', **kwargs)

Compute one or more linear projections of a complex-valued function

Given an input of shape (…, num_rotations, num_sphs), ComplexProjection produces an output of shape (.., num_rotations, num_projections).

Outputs are converted to real numbers by taking the absolute value of the projection output by default, but the real or imaginary components can also be taken instead.

Parameters:

• num_projections – Number of projections (i.e. number of neurons) to create for each rotation
• conversion – Method to make the projection output real: ‘abs’ (absolute value), ‘angle’ (complex phase), ‘real’ (real component), ‘imag’ (imaginary component), or a comma-separated list of these values (i.e. ‘real,imag’)
• activation – Keras activation function for the layer
• kernel_initializer – Keras initializer for the projection weights matrix
• bias_initializer – Keras initializer for the projection bias matrix