nndt package#

Subpackages#

Submodules#

nndt.datagen module#

class nndt.datagen.AbstractDatagen[source]#

Bases: abc.ABC

abstract get(key: jax._src.prng.PRNGKeyArray, epoch: int, **kwargs)[source]#
class nndt.datagen.DataGenForSegmentation(node, cube_spacing=(16, 16, 16), cube_scale=1.0, count=33, step=77, shift_sigma=0.09, scale_range=0.03, rotate_angle=1.0, shift_mul=4, augment=True)[source]#

Bases: nndt.datagen.AbstractDatagen

get(key, epoch, index=None)[source]#
class nndt.datagen.DataGenForShapeRegression(node, spacing=(4.0, 4.0, 4.0), shift_sigma=0.05, augment=True)[source]#

Bases: nndt.datagen.AbstractDatagen

get(key, epoch, **kwargs)[source]#
nndt.datagen._vec_scale_rotate_shift(cube, scale=(1.0, 1.0, 1.0), rotation=(0.0, 0.0, 0.0), shift=(0.0, 0.0, 0.0))#

Vectorized version of _scale_rotate_shift. Takes similar arguments as _scale_rotate_shift but with additional array axes over which _scale_rotate_shift is mapped.

nndt.global_config module#

nndt.global_config.init_code()[source]#

This method rewrites global parameters and allows NNDT works inside the Jupyter notebook environment

Returns

none

nndt.global_config.init_colab(window_size: Tuple[int, int] = (600, 400))[source]#

This method rewrites global parameters and allows NNDT works inside the Google Colaboratory environment

Parameters

window_size – default size of output 3D images

Returns

none

nndt.global_config.init_jupyter()[source]#

This method rewrites global parameters and allows NNDT works inside the jupyter notebook environment

Returns

none

nndt.haiku_modules module#

class nndt.haiku_modules.DescConv(*args, **kwargs)[source]#

Bases: haiku._src.module.Module

Fully convolutional network

class nndt.haiku_modules.LipLinear(*args, **kwargs)[source]#

Bases: haiku._src.module.Module

Layer for implementation of the LipMLP from the article: Liu, Hsueh-Ti Derek, et al. “Learning Smooth Neural Functions via Lipschitz Regularization.” arXiv preprint arXiv:2202.08345 (2022).

get_lipschitz_loss()[source]#
weight_normalization(W, softplus_c)[source]#
class nndt.haiku_modules.LipMLP(*args, **kwargs)[source]#

Bases: haiku._src.module.Module

This is an implementation of the LipMLP from the article: Liu, Hsueh-Ti Derek, et al. “Learning Smooth Neural Functions via Lipschitz Regularization.” arXiv preprint arXiv:2202.08345 (2022).

get_lipschitz_loss()[source]#

nndt.math_core module#

nndt.math_core.barycentric_grid(order: Sequence[Union[int, Sequence[int]]] = (1, - 1), spacing: Sequence[int] = (0, 3), filter_negative: bool = True)[source]#

Analog of nested for cycles in barycentric coordinates. In the 1D case without a free variable, this is linear interpolation. In the 2D case with a free variable, this is the list of ternary plot points. In the ND case, this works like a uniform grid inside N-simplex. If this simplex is defined on the basis of vectors of space.

Parameters

  • order ((Sequence[Union[int, Sequence[int]]], optional)) – Order of iterator in the polynomial (defaults is (1, -1), as for the linear interpolation)

  • spacing ((Sequence[int], optional)) – This is grid spacing for each iterated variable. N-value in some positions is equivalent to jnp.linspace(0,1,N). Zero elements must be zero because this is a technical definition for the free variable.

  • filter_negative ((bool, optional)) – Filter values outside the simple (defaults is True)

Returns

List of vectors inside the simplex. All the vectors have len(spacing) components.

Return type

jnp.ndarray

nndt.math_core.grid_in_cube(spacing=(2, 2, 2), scale=2.0, center_shift=(0.0, 0.0, 0.0)) → jax.Array[source]#

Draw samples from the uniform grid that is defined inside a bounding box with the center in the center_shift and size of scale

Parameters

  • spacing (tuple, optional) – Number of sections along X, Y, and Z axes (default is (2, 2, 2))

  • scale (float, optional) – The scaling factor defines the size of the bounding box (default is 2.)

  • center_shift (tuple, optional) – A tuple of ints of coordinates by which to modify the center of the cube (default is (0., 0., 0.))

Returns

3D mesh-grid with shape (spacing[0], spacing[1], spacing[2], 3)

Return type

ndarray

nndt.math_core.grid_in_cube2(spacing=(4, 4, 4), lower=(- 2, - 2, - 2), upper=(2, 2, 2)) → jax.Array[source]#

Draw samples from the uniform grid that is defined inside a (lower, upper) bounding box

Parameters

  • spacing (tuple, optional) – Number of sections along X, Y, and Z axes (default is (4, 4, 4))

  • lower (tuple, optional) – position of the lower point for the bounding box (default is (-2, -2, -2)

  • upper (tuple, optional) – position of the upper point for the bounding box (default is (2, 2, 2)

Returns

3D mesh-grid with shape (spacing[0], spacing[1], spacing[2], 3)

Return type

ndarray

nndt.math_core.help_barycentric_grid(order: Sequence[Union[int, Sequence[int]]] = (1, - 1))[source]#

Helper for ‘barycentric_grid’ function. This method prints an iteration polynomial for the barycentric coordinates.

Parameters

order (Sequence[Union[int, Sequence[int]]], optional) – Order of iterators (defaults is (1, -1), as for the linear interpolation)

Returns

  • str

  • The text representation of the polynomial

nndt.math_core.rotation_matrix(yaw, pitch, roll)[source]#

Construct rotation matrix from three rotational angle

Parameters

  • yaw – The yaw in radian

  • pitch – The pitch in radian

  • roll – The roll in radian

Returns

nndt.math_core.scale_xyz(xyz, scale=(1.0, 1.0, 1.0))[source]#

Scale array of points to the scale factor.

:param : :type : param xyz: Array of points :param : :type : param scale: The scale factor

Return type

return: Scaled array of points with shape equal to shape of xyz array

nndt.math_core.shift_xyz(xyz, shift=(1.0, 1.0, 1.0))[source]#

Shift array of points to the shift factor.

:param : :type : param xyz: Array of points :param : :type : param scale: The scale factor

Return type

return: Scaled array of points with shape equal to shape of xyz array

nndt.math_core.take_each_n(array: jax.Array, count=1, step=1, shift=0) -> (<class 'jax.Array'>, <class 'jax.Array'>)[source]#

An advanced range iterator that iterates over data and selects elements according to their index. If during iteration the index becomes greater than the array length,

the iteration continues from the beginning of the array.

This function selects elements from an array along the axis zero, which is the first dimension.

Parameters

  • array (ndarray) – The source array

  • count (int, optional) – The number of elements to take (default is 1)

  • step (int, optional) – The step of iterator

  • shift (int, optional) – Index shift for the first index (default is 0)

Returns

an array of indices of the elements taken from the source array an array of elements from the source array corresponding to the selected indices

Return type

(ndarray, ndarray)

nndt.math_core.train_test_split(array: jax.Array, rng: jax._src.prng.PRNGKeyArray, test_size: float = 0.3) -> (<class 'list'>, <class 'list'>)[source]#

Split array to test and train subset. This is analog of model_selection.train_test_split in sklearn.

Parameters

  • array (jnp.ndarray :) – Array for split

  • rng (KeyArray :) – Jax key for a random generator

  • test_size (float:) – Percent of test subset in the array

Returns

List of indexes for test and train subsets

Return type

(list, list)

nndt.math_core.uniform_in_cube(rng_key: jax._src.prng.PRNGKeyArray, count=100, lower=(- 2, - 2, - 2), upper=(2, 2, 2)) → jax.Array[source]#
Parameters

  • rng_key (KeyArray) – Jax key for a random generator

  • count (int, optional) – Size of sampling (default is 100)

  • lower (tuple, optional) – position of the lower point for the bounding box (default is (-2, -2, -2)

  • upper (tuple, optional) – position of the upper point for the bounding box (default is (2, 2, 2)

Returns

  • ndarray

  • An array of random points (shape is (count, 3))

nndt.primitive_sdf module#

class nndt.primitive_sdf.AbstractSDF[source]#

Bases: object

abstract property bbox: ((<class 'float'>, <class 'float'>, <class 'float'>), (<class 'float'>, <class 'float'>, <class 'float'>))#

Return the minimal bounding box around the implicitly defined object. :return: (X_min, Y_min, Z_min) , (X_max, Y_max, Z_max)

property fun: Callable#

Get the SDF function in scalar form :return: f(x,y,z) = distance

request(ps_xyz: jax.Array) → jax.Array[source]#

Get SDF values for the requested location on the physical space. :return: distance values

property vec_fun: Callable#

Get the SDF function in vector form. Vectorization is performed along the zero axis. :return: f(vec_x, vec_y, vec_z) = vec_distance

property vec_fun_dx: Callable#

Get the gradient of the SDF function over the X-axis. Vectorization is performed along the zero axis. :return: df/dx(vec_x, vec_y, vec_z)

property vec_fun_dy: Callable#

Get the gradient of the SDF function over the Y-axis. Vectorization is performed along the zero axis. :return: df/dy(vec_x, vec_y, vec_z)

property vec_fun_dz: Callable#

Get the gradient of the SDF function over the Z-axis. Vectorization is performed along the zero axis. :return: df/dz(vec_x, vec_y, vec_z)

class nndt.primitive_sdf.SphereSDF(center=(0.0, 0.0, 0.0), radius=1.0)[source]#

Bases: nndt.primitive_sdf.AbstractSDF

This is a sphere geometrical primitive.

__init__(center=(0.0, 0.0, 0.0), radius=1.0)[source]#

This is a sphere geometrical primitive.

Parameters

  • center – center of the sphere

  • radius – radius of the sphere

property bbox: ((<class 'float'>, <class 'float'>, <class 'float'>), (<class 'float'>, <class 'float'>, <class 'float'>))#

Return the minimal bounding box around the implicitly defined object. :return: (X_min, Y_min, Z_min) , (X_max, Y_max, Z_max)

nndt.primitive_sdf.fun2vec_and_grad(prim)[source]#

nndt.trainable_task module#

class nndt.trainable_task.AbstractTrainableTask[source]#

Bases: object

abstract init_and_functions(rng_key: jax._src.prng.PRNGKeyArray) -> (<function namedtuple at 0x7f673132d5e0>, <function namedtuple at 0x7f673132d5e0>)[source]#
abstract init_data() → collections.namedtuple[source]#
class nndt.trainable_task.ApproximateSDF(mlp_layers=(64, 64, 64, 64, 64, 64, 64, 64, 1), batch_size=262144, model_number=2)[source]#

Bases: nndt.trainable_task.AbstractTrainableTask

This is a trainable task for the problem of shape interpolation between several 3d objects. This class employs the usual multi-layer perceptron.

class DATA(X, Y, Z, T, P, SDF)[source]#

Bases: NamedTuple

P: jax.Array#

Alias for field number 4

SDF: jax.Array#

Alias for field number 5

T: jax.Array#

Alias for field number 3

X: jax.Array#

Alias for field number 0

Y: jax.Array#

Alias for field number 1

Z: jax.Array#

Alias for field number 2

FUNC#

alias of nndt.trainable_task.ApproximateSDF_DATA

init_and_functions(rng_key: jax._src.prng.PRNGKeyArray) -> (<function namedtuple at 0x7f673132d5e0>, <function namedtuple at 0x7f673132d5e0>)[source]#
init_data() → collections.namedtuple[source]#
class nndt.trainable_task.ApproximateSDFLipMLP(mlp_layers=(64, 64, 64, 64, 64, 64, 64, 64, 1), batch_size=262144, model_number=2, lip_alpha=1e-06, negative_beta=0.0)[source]#

Bases: nndt.trainable_task.AbstractTrainableTask

This is a trainable task for the problem of shape interpolation between several 3d objects. This class employs a multi-layer perceptron with Lipschitz regularization (LipMLP).

class DATA(X, Y, Z, T, P, SDF)[source]#

Bases: NamedTuple

P: jax.Array#

Alias for field number 4

SDF: jax.Array#

Alias for field number 5

T: jax.Array#

Alias for field number 3

X: jax.Array#

Alias for field number 0

Y: jax.Array#

Alias for field number 1

Z: jax.Array#

Alias for field number 2

FUNC#

alias of nndt.trainable_task.ApproximateSDFLipMLP_DATA

init_and_functions(rng_key: jax._src.prng.PRNGKeyArray) -> (<function namedtuple at 0x7f673132d5e0>, <function namedtuple at 0x7f673132d5e0>)[source]#
init_data() → collections.namedtuple[source]#
class nndt.trainable_task.ApproximateSDFLipMLP2(mlp_layers=(64, 64, 64, 64, 64, 64, 64, 64, 1), batch_size=262144, model_number=2, lip_alpha=1e-06)[source]#

Bases: nndt.trainable_task.AbstractTrainableTask

This is a trainable task for the problem of shape interpolation between several 3d objects. This class employs a multi-layer perceptron with Lipschitz regularization (LipMLP).

class DATA(X, Y, Z, T, P, SDF, WEIGHT)[source]#

Bases: NamedTuple

P: jax.Array#

Alias for field number 4

SDF: jax.Array#

Alias for field number 5

T: jax.Array#

Alias for field number 3

WEIGHT: jax.Array#

Alias for field number 6

X: jax.Array#

Alias for field number 0

Y: jax.Array#

Alias for field number 1

Z: jax.Array#

Alias for field number 2

FUNC#

alias of nndt.trainable_task.ApproximateSDFLipMLP2_DATA

init_and_functions(rng_key: jax._src.prng.PRNGKeyArray) -> (<function namedtuple at 0x7f673132d5e0>, <function namedtuple at 0x7f673132d5e0>)[source]#
init_data() → collections.namedtuple[source]#
class nndt.trainable_task.Eikonal3D(fun_sdf_domain: Callable[[float, float, float], float], fun_sdf_start: Callable[[float, float, float], float], mlp_layers=(16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 1), lambda_grad=0.1, lambda_non_negativity=0.1, batch_size=100)[source]#

Bases: nndt.trainable_task.AbstractTrainableTask

This is a solver for the Eikonal equation.

class DATA(X, Y, Z)[source]#

Bases: NamedTuple

X: jax.Array#

Alias for field number 0

Y: jax.Array#

Alias for field number 1

Z: jax.Array#

Alias for field number 2

class FUNC(nn, nn_dx, nn_dy, nn_dz, main_loss)[source]#

Bases: NamedTuple

main_loss: Callable[[float, float, float], float]#

Alias for field number 4

nn: Callable[[float, float, float], float]#

Alias for field number 0

nn_dx: Callable[[float, float, float], float]#

Alias for field number 1

nn_dy: Callable[[float, float, float], float]#

Alias for field number 2

nn_dz: Callable[[float, float, float], float]#

Alias for field number 3

init_and_functions(rng_key: jax._src.prng.PRNGKeyArray) -> (<function namedtuple at 0x7f673132d5e0>, <function namedtuple at 0x7f673132d5e0>)[source]#
init_data() → collections.namedtuple[source]#
class nndt.trainable_task.SimpleSDF(mlp_layers=(32, 32, 32, 32, 32, 32, 32, 32, 1), batch_size=262144)[source]#

Bases: nndt.trainable_task.AbstractTrainableTask

This is a trainable task for the problem of shape interpolation for one 3d object. This class employs the usual multi-layer perceptron.

class DATA(X, Y, Z, SDF)[source]#

Bases: NamedTuple

SDF: jax.Array#

Alias for field number 3

X: jax.Array#

Alias for field number 0

Y: jax.Array#

Alias for field number 1

Z: jax.Array#

Alias for field number 2

class FUNC(sdf, vec_sdf, sdf_dx, sdf_dy, sdf_dz, vec_sdf_dx, vec_sdf_dy, vec_sdf_dz, vec_main_loss)[source]#

Bases: NamedTuple

sdf: Callable#

Alias for field number 0

sdf_dx: Callable#

Alias for field number 2

sdf_dy: Callable#

Alias for field number 3

sdf_dz: Callable#

Alias for field number 4

vec_main_loss: Callable#

Alias for field number 8

vec_sdf: Callable#

Alias for field number 1

vec_sdf_dx: Callable#

Alias for field number 5

vec_sdf_dy: Callable#

Alias for field number 6

vec_sdf_dz: Callable#

Alias for field number 7

init_and_functions(rng_key: jax._src.prng.PRNGKeyArray) -> (<function namedtuple at 0x7f673132d5e0>, <function namedtuple at 0x7f673132d5e0>)[source]#
class nndt.trainable_task.SurfaceSegmentation(spacing=(16, 16, 16), conv_kernel=32, conv_depth=4, num_classes=3, batch_size=128)[source]#

Bases: nndt.trainable_task.AbstractTrainableTask

This is a trainable task for the problem of supervised surface segmentation. This class employs the fully-convolutional neural network.

class DATA(SDF_CUBE, CLASS)#

Bases: tuple

CLASS#

Alias for field number 1

SDF_CUBE#

Alias for field number 0

class FUNC(nn, main_loss, metric_accuracy)#

Bases: tuple

main_loss#

Alias for field number 1

metric_accuracy#

Alias for field number 2

nn#

Alias for field number 0

init_and_functions(rng_key: jax._src.prng.PRNGKeyArray) -> (<function namedtuple at 0x7f673132d5e0>, <function namedtuple at 0x7f673132d5e0>)[source]#
init_data() → collections.namedtuple[source]#

nndt.vizualize module#

class nndt.vizualize.BasicVizualization(folder: str, experiment_name: Optional[str] = None, print_on_each_epoch=20)[source]#

Bases: object

Simple MLOps class for storing the train history and visualization of intermediate results

__init__(folder: str, experiment_name: Optional[str] = None, print_on_each_epoch=20)[source]#

Simple MLOps class for storing the train history and visualization of intermediate results

Parameters

  • folder – folder for store results

  • experiment_name – name for an experiments

  • print_on_each_epoch – this parameter helps to control intermediate result output

draw_loss(name, history)[source]#

Save the training history in .jpg

Parameters

  • name (string) – File name

  • (_type_) (history) – List of loss values over epochs

is_print_on_epoch(epoch)[source]#

Check if this is the right epoch to print results

Parameters

epoch (int) – epoch number

Returns

Should we print on this step?

Return type

bool

iter(epoch_num)[source]#

Return iterators for the main train cycle

Parameters

epoch_num (int) – number of epoch

Return type

instance of IteratorWithTimeMeasurements

record(dict)[source]#
save_3D_array(name, array, section_img=True)[source]#

Save the 3D array to a file and section of this array as images

Parameters

  • name (string) – File name

  • array (array) – 3D array to save

  • section_img (bool) – If true, this saves three plane section of 3D array (defaults to True)

save_mesh(name, save_method, dict_)[source]#

Save mesh to .vtp file with data

Parameters

  • name (string) – filename

  • save_method (SaveMesh) – SaveMesh instance from NNDT space (v0.0.1 or v0.0.2)

  • dict (dict) – name_value

save_state(name, state)[source]#

Save the neural network state into the file

Parameters

  • name (string) – File name

  • state ((_type_)) – The state to save

save_txt(name, summary)[source]#

Save string data to .txt file

Parameters

  • name (string) – File name

  • summary (string) – The text to save

sdt_to_obj(filename: str, array: Union[jax.Array, numpy.ndarray], level: float = 0.0)[source]#

Run marching cubes over SDT and save results to file

Parameters

  • filename (string) – File name

  • array (ndarray) – Signed distance tensor (SDT)

  • level (float) – Isosurface level (defaults to 0.).

class nndt.vizualize.IteratorWithTimeMeasurements(basic_viz, epochs)[source]#

Bases: object

Iterator that records and prints the epoch number and time spent from the start of iterations

nndt.vizualize.save_3D_slices(array: Union[numpy.ndarray, jax.Array], path: Optional[str] = None, slice_num: int = 5, include_boundary=True, figsize=None, levels=(0.0,), level_colors=('white',), **kwargs)[source]#

Generates a panel of images with slices of the 3D array. This is a helper function for studying 3D tensors.

Parameters

  • path – path to the image for write

  • array – studied array

  • slice_num – number of slices over array axis

  • include_boundary – If True, the image will include boundaries of the array with indexes 0 and len(array)-1

  • figsize – the size of the image. If None, the size will be calculated according to the number of panels.

  • levels – Isoline values. This param is ignored if RGB/RGBA image is passed.

  • level_colors – Isoline colors. This param is ignored if RGB/RGBA image is passed.

  • kwargs – parameter set that is passed to the .imshow() method

Returns

none

nndt.vizualize.save_sdt_as_obj(array: Union[jax.Array, numpy.ndarray], path: str, level: float = 0.0)[source]#

Run marching cubes over SDT and save results to file

Parameters

  • filename (string) – File name

  • array (ndarray) – Signed distance tensor (SDT)

  • level (float) – Isosurface level (defaults to 0.).

Module contents#