diff --git a/python/rascaline/rascaline/utils/clebsch_gordan/__init__.py b/python/rascaline/rascaline/utils/clebsch_gordan/__init__.py index 38777d18e..62a140ae5 100644 --- a/python/rascaline/rascaline/utils/clebsch_gordan/__init__.py +++ b/python/rascaline/rascaline/utils/clebsch_gordan/__init__.py @@ -1,7 +1,10 @@ -from .clebsch_gordan import correlate_density, correlate_density_metadata # noqa +from .correlate_density import correlate_density, correlate_density_metadata # noqa +from .correlate_tensors import correlate_tensors, correlate_tensors_metadata # noqa __all__ = [ "correlate_density", "correlate_density_metadata", + "correlate_tensors", + "correlate_tensors_metadata", ] diff --git a/python/rascaline/rascaline/utils/clebsch_gordan/clebsch_gordan.py b/python/rascaline/rascaline/utils/clebsch_gordan/_clebsch_gordan.py similarity index 64% rename from python/rascaline/rascaline/utils/clebsch_gordan/clebsch_gordan.py rename to python/rascaline/rascaline/utils/clebsch_gordan/_clebsch_gordan.py index 5572c7598..e2c18914e 100644 --- a/python/rascaline/rascaline/utils/clebsch_gordan/clebsch_gordan.py +++ b/python/rascaline/rascaline/utils/clebsch_gordan/_clebsch_gordan.py @@ -1,5 +1,7 @@ """ -Module for computing Clebsch-gordan iterations with metatensor TensorMaps. +Private module containing helper functions for public modules +:py:mod:`correlate_tensors` and :py:mod:`correlate_density` that compute +Clebsch-gordan tensor products on metatensor :py:class:`TensorMap` objects. """ import itertools from typing import List, Optional, Tuple, Union @@ -9,255 +11,6 @@ from . import _cg_cache, _dispatch -# ====================================================================== -# ===== Public API functions -# ====================================================================== - - -def correlate_density( - density: TensorMap, - correlation_order: int, - angular_cutoff: Optional[int] = None, - selected_keys: Optional[Union[Labels, List[Labels]]] = None, - skip_redundant: Optional[Union[bool, List[bool]]] = False, - output_selection: Optional[Union[bool, List[bool]]] = None, -) -> Union[TensorMap, List[TensorMap]]: - """ - Takes iterative Clebsch-Gordan (CG) tensor products of a density descriptor - with itself up to the desired correlation order. Returns - :py:class:`TensorMap`(s) corresponding to the density correlations output - from the specified iteration(s). - - A density descriptor necessarily is body order 2 (i.e. correlation order 1), - but can be single- or multi-center. The output is a :py:class:`list` of - density correlations for each iteration specified in `output_selection`, up - to the target order passed in `correlation_order`. By default only the last - correlation (i.e. the correlation of order ``correlation_order``) is - returned. - - This function is an iterative special case of the more general - :py:func:`correlate_tensors`. As a density is being correlated with itself, - some redundant CG tensor products can be skipped with the `skip_redundant` - keyword. - - Selections on the angular and parity channels at each iteration can also be - controlled with arguments `angular_cutoff`, `angular_selection` and - `parity_selection`. - - :param density: A density descriptor of body order 2 (correlation order 1), - in :py:class:`TensorMap` format. This may be, for example, a rascaline - :py:class:`SphericalExpansion` or :py:class:`LodeSphericalExpansion`. - Alternatively, this could be multi-center descriptor, such as a pair - density. - :param correlation_order: The desired correlation order of the output - descriptor. Must be >= 1. - :param angular_cutoff: The maximum angular channel to compute at any given - CG iteration, applied globally to all iterations until the target - correlation order is reached. - :param selected_keys: :py:class:`Labels` or `List[:py:class:`Labels`]` - specifying the angular and/or parity channels to output at each - iteration. All :py:class:`Labels` objects passed here must only contain - key names "spherical_harmonics_l" and "inversion_sigma". If a single - :py:class:`Labels` object is passed, this is applied to the final - iteration only. If a :py:class:`list` of :py:class:`Labels` objects is - passed, each is applied to its corresponding iteration. If None is - passed, all angular and parity channels are output at each iteration, - with the global `angular_cutoff` applied if specified. - :param skip_redundant: Whether to skip redundant CG combinations. Defaults - to False, which means all combinations are performed. If a - :py:class:`list` of :py:class:`bool` is passed, this is applied to each - iteration. If a single :py:class:`bool` is passed, this is applied to - all iterations. - :param output_selection: A :py:class:`list` of :py:class:`bool` specifying - whether to output a :py:class:`TensorMap` for each iteration. If a - single :py:class:`bool` is passed as True, outputs from all iterations - will be returned. If a :py:class:`list` of :py:class:`bool` is passed, - this controls the output at each corresponding iteration. If None is - passed, only the final iteration is output. - - :return: A :py:class:`list` of :py:class:`TensorMap` corresponding to the - density correlations output from the specified iterations. If the output - from a single iteration is requested, a :py:class:`TensorMap` is - returned instead. - """ - return _correlate_density( - density, - correlation_order, - angular_cutoff, - selected_keys, - skip_redundant, - output_selection, - compute_metadata_only=False, - sparse=True, # sparse CG cache by default - ) - - -def correlate_density_metadata( - density: TensorMap, - correlation_order: int, - angular_cutoff: Optional[int] = None, - selected_keys: Optional[Union[Labels, List[Labels]]] = None, - skip_redundant: Optional[Union[bool, List[bool]]] = False, - output_selection: Optional[Union[bool, List[bool]]] = None, -) -> Union[TensorMap, List[TensorMap]]: - """ - Returns the metadata-only :py:class:`TensorMap`(s) that would be output by - the function :py:func:`correlate_density` under the same settings, without - perfoming the actual Clebsch-Gordan tensor products. See this function for - full documentation. - """ - - return _correlate_density( - density, - correlation_order, - angular_cutoff, - selected_keys, - skip_redundant, - output_selection, - compute_metadata_only=True, - ) - - -# ==================================================================== -# ===== Private functions that do the work on the TensorMap level -# ==================================================================== - - -def _correlate_density( - density: TensorMap, - correlation_order: int, - angular_cutoff: Optional[int] = None, - selected_keys: Optional[Union[Labels, List[Labels]]] = None, - skip_redundant: Optional[Union[bool, List[bool]]] = False, - output_selection: Optional[Union[bool, List[bool]]] = None, - compute_metadata_only: bool = False, - sparse: bool = True, -) -> Union[TensorMap, List[TensorMap]]: - """ - Performs the density correlations for public functions - :py:func:`correlate_density` and :py:func:`correlate_density_metadata`. - """ - # Check inputs - if correlation_order <= 1: - raise ValueError("`correlation_order` must be > 1") - # TODO: implement combinations of gradients too - if _dispatch.any([len(list(block.gradients())) > 0 for block in density]): - raise NotImplementedError( - "Clebsch Gordan combinations with gradients not yet implemented." - " Use metatensor.remove_gradients to remove gradients from the input." - ) - # Check metadata - if not ( - _dispatch.all(density.keys.names == ["spherical_harmonics_l", "species_center"]) - or _dispatch.all( - density.keys.names - == ["spherical_harmonics_l", "species_center", "species_neighbor"] - ) - ): - raise ValueError( - "input `density` must have key names" - ' ["spherical_harmonics_l", "species_center"] or' - ' ["spherical_harmonics_l", "species_center", "species_neighbor"]' - ) - if not _dispatch.all(density.component_names == ["spherical_harmonics_m"]): - raise ValueError( - "input `density` must have a single component" - " axis with name `spherical_harmonics_m`" - ) - n_iterations = correlation_order - 1 # num iterations - density = _standardize_keys(density) # standardize metadata - density_correlation = density # create a copy to combine with itself - - # Parse the selected keys - selected_keys = _parse_selected_keys( - n_iterations=n_iterations, - angular_cutoff=angular_cutoff, - selected_keys=selected_keys, - like=density.keys.values, - ) - # Parse the bool flags that control skipping of redundant CG combinations - # and TensorMap output from each iteration - skip_redundant, output_selection = _parse_bool_iteration_filters( - n_iterations, - skip_redundant=skip_redundant, - output_selection=output_selection, - ) - - # Pre-compute the keys needed to perform each CG iteration - key_metadata = _precompute_keys( - density.keys, - density.keys, - n_iterations=n_iterations, - selected_keys=selected_keys, - skip_redundant=skip_redundant, - ) - # Compute CG coefficient cache - if compute_metadata_only: - cg_cache = None - else: - angular_max = max( - _dispatch.concatenate( - [density.keys.column("spherical_harmonics_l")] - + [mdata[2].column("spherical_harmonics_l") for mdata in key_metadata] - ) - ) - # TODO: keys have been precomputed, so perhaps we don't need to - # compute all CG coefficients up to angular_max here. - # TODO: use sparse cache by default until we understand under which - # circumstances (and if) dense is faster. - cg_cache = _cg_cache.ClebschGordanReal(angular_max, sparse=sparse) - - # Perform iterative CG tensor products - density_correlations = [] - for iteration in range(n_iterations): - # Define the correlation order of the current iteration - correlation_order_it = iteration + 2 - - # Combine block pairs - blocks_out = [] - for key_1, key_2, key_out in zip(*key_metadata[iteration]): - block_out = _combine_blocks_same_samples( - density_correlation[key_1], - density[key_2], - key_out["spherical_harmonics_l"], - cg_cache, - compute_metadata_only=compute_metadata_only, - ) - blocks_out.append(block_out) - keys_out = key_metadata[iteration][2] - density_correlation = TensorMap(keys=keys_out, blocks=blocks_out) - - # If this tensor is to be included in the output, move the [l1, l2, ...] - # keys to properties and store - if output_selection[iteration]: - density_correlations.append( - density_correlation.keys_to_properties( - [f"l{i}" for i in range(1, correlation_order_it + 1)] - + [f"k{i}" for i in range(2, correlation_order_it)] - ) - ) - - # Drop redundant key names. TODO: these should be part of the global - # matadata associated with the TensorMap. Awaiting this functionality in - # metatensor. - for i, tensor in enumerate(density_correlations): - keys = tensor.keys - if len(_dispatch.unique(tensor.keys.column("order_nu"))) == 1: - keys = keys.remove(name="order_nu") - if len(_dispatch.unique(tensor.keys.column("inversion_sigma"))) == 1: - keys = keys.remove(name="inversion_sigma") - density_correlations[i] = TensorMap( - keys=keys, blocks=[b.copy() for b in tensor.blocks()] - ) - - # Return a single TensorMap in the simple case - if len(density_correlations) == 1: - return density_correlations[0] - - # Otherwise return a list of TensorMaps - return density_correlations - - # ================================================================== # ===== Functions to handle metadata # ================================================================== @@ -445,7 +198,9 @@ def _precompute_keys( `selected_keys`. If `skip_redundant` is True, then keys that represent redundant CG - operations are not included in the output keys at each step. + operations are not included in the output keys at each step. This is only + valid when iteratively correlating a density (i.e. correlation order 1) + descriptor with itself. """ keys_metadata = [] keys_out = keys_1 @@ -464,6 +219,9 @@ def _precompute_keys( ) if skip_redundant[iteration]: + # Removing redundant keys is only valid when correlating a density + # iteratively with itself + assert set(keys_2.column("order_nu")) == {1} keys_1_entries, keys_2_entries, keys_out = _remove_redundant_keys( keys_1_entries, keys_2_entries, keys_out ) @@ -525,19 +283,41 @@ def _precompute_keys_full_product( the pair of blocks indexed by correspoding key pairs in the first two lists. """ # Get the correlation order of the first TensorMap. - unique_nu = _dispatch.unique(keys_1.column("order_nu")) - if len(unique_nu) > 1: + unique_nu1 = _dispatch.unique(keys_1.column("order_nu")) + if len(unique_nu1) > 1: raise ValueError( "keys_1 must correspond to a tensor of a single correlation order." - f" Found {len(unique_nu)} body orders: {unique_nu}" + f" Found {len(unique_nu1)} body orders: {unique_nu1}" ) - nu1 = unique_nu[0] + nu1 = unique_nu1[0] + + # Get the correlation order of the second TensorMap. + unique_nu2 = _dispatch.unique(keys_2.column("order_nu")) + if len(unique_nu2) > 1: + raise ValueError( + "keys_1 must correspond to a tensor of a single correlation order." + f" Found {len(unique_nu2)} body orders: {unique_nu2}" + ) + nu2 = unique_nu2[0] # Define new correlation order of output TensorMap - nu = nu1 + 1 + nu = nu1 + nu2 # The correlation order of the second TensorMap should be nu = 1. - assert _dispatch.all(keys_2.column("order_nu") == 1) + assert nu2 == 1 + + # Define the sub-parts of the required key names. These are: + # 1) The base key names both keys should have: correlation order, parity, + # lambda + base_key_names = ["order_nu", "inversion_sigma", "spherical_harmonics_l"] + # 2) The list of l values, depending on the correlation order + l_list_names_1 = [f"l{angular_l}" for angular_l in range(1, nu1 + 1)] + l_list_names_2 = [f"l{angular_l}" for angular_l in range(1, nu2 + 1)] + # 3) The list of k values, depending on the correlation order + k_list_names_1 = [f"k{k}" for k in range(2, nu1)] + k_list_names_2 = [f"k{k}" for k in range(2, nu2)] + # 4) Any other sparse keys - such as species center, species neighbor, etc. + other_key_names = # If nu1 = 1, the key names don't yet have any "lx" columns if nu1 == 1: @@ -548,6 +328,7 @@ def _precompute_keys_full_product( new_l_list_names = l_list_names + [f"l{nu}"] # Check key names + assert _dispatch.all( keys_1.names == ["order_nu", "inversion_sigma", "spherical_harmonics_l", "species_center"] @@ -668,6 +449,9 @@ def _remove_redundant_keys( # Get and check the correlation order of the input keys nu1 = keys_1_entries[0]["order_nu"] nu2 = keys_2_entries[0]["order_nu"] + + # Removing redundant keys is only valid when correlating a density + # iteratively with itself assert nu2 == 1 # Get the correlation order of the output TensorMap diff --git a/python/rascaline/rascaline/utils/clebsch_gordan/correlate_density.py b/python/rascaline/rascaline/utils/clebsch_gordan/correlate_density.py new file mode 100644 index 000000000..9704548a9 --- /dev/null +++ b/python/rascaline/rascaline/utils/clebsch_gordan/correlate_density.py @@ -0,0 +1,246 @@ +""" +Module for computing Clebsch-gordan tensor product iterations on density (i.e. +correlation order 1) tensors in TensorMap form, where the samples are +equivalent. + +This is a special (and iterative) case of the more general +:py:mod:`correlate_tensors` module. +""" +from typing import List, Optional, Union + +from metatensor import Labels, TensorMap + +from . import _cg_cache, _clebsch_gordan, _dispatch + + +# ====================================================================== +# ===== Public API functions +# ====================================================================== + + +def correlate_density( + density: TensorMap, + correlation_order: int, + angular_cutoff: Optional[int] = None, + selected_keys: Optional[Union[Labels, List[Labels]]] = None, + skip_redundant: Optional[Union[bool, List[bool]]] = False, + output_selection: Optional[Union[bool, List[bool]]] = None, +) -> Union[TensorMap, List[TensorMap]]: + """ + Takes iterative Clebsch-Gordan (CG) tensor products of a density descriptor + with itself up to the desired correlation order. Returns + :py:class:`TensorMap`(s) corresponding to the density correlations output + from the specified iteration(s). + + A density descriptor necessarily is body order 2 (i.e. correlation order 1), + but can be single- or multi-center. The output is a :py:class:`list` of + density correlations for each iteration specified in `output_selection`, up + to the target order passed in `correlation_order`. By default only the last + correlation (i.e. the correlation of order ``correlation_order``) is + returned. + + This function is an iterative special case of the more general + :py:mod:`correlate_tensors`. As a density is being correlated with itself, + some redundant CG tensor products can be skipped with the `skip_redundant` + keyword. + + Selections on the angular and parity channels at each iteration can also be + controlled with arguments `selected_keys`. + + :param density: A density descriptor of body order 2 (correlation order 1), + in :py:class:`TensorMap` format. This may be, for example, a rascaline + :py:class:`SphericalExpansion` or :py:class:`LodeSphericalExpansion`. + Alternatively, this could be multi-center descriptor, such as a pair + density. + :param correlation_order: The desired correlation order of the output + descriptor. Must be >= 1. + :param angular_cutoff: The maximum angular channel to compute at any given + CG iteration, applied globally to all iterations until the target + correlation order is reached. + :param selected_keys: :py:class:`Labels` or `List[:py:class:`Labels`]` + specifying the angular and/or parity channels to output at each + iteration. All :py:class:`Labels` objects passed here must only contain + key names "spherical_harmonics_l" and "inversion_sigma". If a single + :py:class:`Labels` object is passed, this is applied to the final + iteration only. If a :py:class:`list` of :py:class:`Labels` objects is + passed, each is applied to its corresponding iteration. If None is + passed, all angular and parity channels are output at each iteration, + with the global `angular_cutoff` applied if specified. + :param skip_redundant: Whether to skip redundant CG combinations. Defaults + to False, which means all combinations are performed. If a + :py:class:`list` of :py:class:`bool` is passed, this is applied to each + iteration. If a single :py:class:`bool` is passed, this is applied to + all iterations. + :param output_selection: A :py:class:`list` of :py:class:`bool` specifying + whether to output a :py:class:`TensorMap` for each iteration. If a + single :py:class:`bool` is passed as True, outputs from all iterations + will be returned. If a :py:class:`list` of :py:class:`bool` is passed, + this controls the output at each corresponding iteration. If None is + passed, only the final iteration is output. + + :return: A :py:class:`list` of :py:class:`TensorMap` corresponding to the + density correlations output from the specified iterations. If the output + from a single iteration is requested, a :py:class:`TensorMap` is + returned instead. + """ + return _correlate_density( + density, + correlation_order, + angular_cutoff, + selected_keys, + skip_redundant, + output_selection, + compute_metadata_only=False, + sparse=True, # sparse CG cache by default + ) + + +def correlate_density_metadata( + density: TensorMap, + correlation_order: int, + angular_cutoff: Optional[int] = None, + selected_keys: Optional[Union[Labels, List[Labels]]] = None, + skip_redundant: Optional[Union[bool, List[bool]]] = False, + output_selection: Optional[Union[bool, List[bool]]] = None, +) -> Union[TensorMap, List[TensorMap]]: + """ + Returns the metadata-only :py:class:`TensorMap`(s) that would be output by + the function :py:func:`correlate_density` under the same settings, without + perfoming the actual Clebsch-Gordan tensor products. See this function for + full documentation. + """ + + return _correlate_density( + density, + correlation_order, + angular_cutoff, + selected_keys, + skip_redundant, + output_selection, + compute_metadata_only=True, + ) + + +# ==================================================================== +# ===== Private functions that do the work on the TensorMap level +# ==================================================================== + + +def _correlate_density( + density: TensorMap, + correlation_order: int, + angular_cutoff: Optional[int] = None, + selected_keys: Optional[Union[Labels, List[Labels]]] = None, + skip_redundant: Optional[Union[bool, List[bool]]] = False, + output_selection: Optional[Union[bool, List[bool]]] = None, + compute_metadata_only: bool = False, + sparse: bool = True, +) -> Union[TensorMap, List[TensorMap]]: + """ + Performs the density correlations for public functions + :py:func:`correlate_density` and :py:func:`correlate_density_metadata`. + """ + # Check inputs + if correlation_order <= 1: + raise ValueError("`correlation_order` must be > 1") + # TODO: implement combinations of gradients too + if _dispatch.any([len(list(block.gradients())) > 0 for block in density]): + raise NotImplementedError( + "Clebsch Gordan combinations with gradients not yet implemented." + " Use metatensor.remove_gradients to remove gradients from the input." + ) + # Check metadata + if not ( + _dispatch.all(density.keys.names == ["spherical_harmonics_l", "species_center"]) + or _dispatch.all( + density.keys.names + == ["spherical_harmonics_l", "species_center", "species_neighbor"] + ) + ): + raise ValueError( + "input `density` must have key names" + ' ["spherical_harmonics_l", "species_center"] or' + ' ["spherical_harmonics_l", "species_center", "species_neighbor"]' + ) + if not _dispatch.all(density.component_names == ["spherical_harmonics_m"]): + raise ValueError( + "input `density` must have a single component" + " axis with name `spherical_harmonics_m`" + ) + n_iterations = correlation_order - 1 # num iterations + density = _clebsch_gordan._standardize_keys(density) # standardize metadata + density_correlation = density # create a copy to combine with itself + + # Parse the selected keys + selected_keys = _clebsch_gordan._parse_selected_keys( + n_iterations=n_iterations, + angular_cutoff=angular_cutoff, + selected_keys=selected_keys, + like=density.keys.values, + ) + # Parse the bool flags that control skipping of redundant CG combinations + # and TensorMap output from each iteration + skip_redundant, output_selection = _clebsch_gordan._parse_bool_iteration_filters( + n_iterations, + skip_redundant=skip_redundant, + output_selection=output_selection, + ) + + # Pre-compute the keys needed to perform each CG iteration + key_metadata = _clebsch_gordan._precompute_keys( + density.keys, + density.keys, + n_iterations=n_iterations, + selected_keys=selected_keys, + skip_redundant=skip_redundant, + ) + # Compute CG coefficient cache + if compute_metadata_only: + cg_cache = None + else: + angular_max = max( + _dispatch.concatenate( + [density.keys.column("spherical_harmonics_l")] + + [mdata[2].column("spherical_harmonics_l") for mdata in key_metadata] + ) + ) + # TODO: keys have been precomputed, so perhaps we don't need to + # compute all CG coefficients up to angular_max here. + # TODO: use sparse cache by default until we understand under which + # circumstances (and if) dense is faster. + cg_cache = _cg_cache.ClebschGordanReal(angular_max, sparse=sparse) + + # Perform iterative CG tensor products + density_correlations = [] + for iteration in range(n_iterations): + # Define the correlation order of the current iteration + correlation_order_it = iteration + 2 + + # Combine block pairs + blocks_out = [] + for key_1, key_2, key_out in zip(*key_metadata[iteration]): + block_out = _clebsch_gordan._combine_blocks_same_samples( + density_correlation[key_1], + density[key_2], + key_out["spherical_harmonics_l"], + cg_cache, + compute_metadata_only=compute_metadata_only, + ) + blocks_out.append(block_out) + keys_out = key_metadata[iteration][2] + density_correlation = TensorMap(keys=keys_out, blocks=blocks_out) + + # If this tensor is to be included in the output, move the [l1, l2, ...] + # keys to properties and store + if output_selection[iteration]: + density_correlations.append( + density_correlation.keys_to_properties( + [f"l{i}" for i in range(1, correlation_order_it + 1)] + + [f"k{i}" for i in range(2, correlation_order_it)] + ) + ) + + if len(density_correlations) == 1: # simple case: single TensorMap + return density_correlations[0] + + return density_correlations # list of TensorMaps diff --git a/python/rascaline/rascaline/utils/clebsch_gordan/correlate_tensors.py b/python/rascaline/rascaline/utils/clebsch_gordan/correlate_tensors.py new file mode 100644 index 000000000..d34b3d8a8 --- /dev/null +++ b/python/rascaline/rascaline/utils/clebsch_gordan/correlate_tensors.py @@ -0,0 +1,174 @@ +""" +Module for computing the Clebsch-gordan tensor product on arbitrary tensors in +TensorMap form, where the samples are equivalent. + +For a special iterative case specifically for densities, see the +:py:mod:`correlate_density` module. +""" +from typing import List, Optional, Union + +from metatensor import Labels, TensorMap + +from . import _cg_cache, _clebsch_gordan, _dispatch + + +# ====================================================================== +# ===== Public API functions +# ====================================================================== + + +def correlate_tensors( + tensor_1: TensorMap, + tensor_2: TensorMap, + selected_keys: Optional[Union[Labels, List[Labels]]] = None, +) -> TensorMap: + """ + Takes a Clebsch-Gordan (CG) tensor product of two tensors in + :py:class:`TensorMap` form, and returns the result in :py:class:`TensorMap`. + + The keys of `tensor_1` and `tensor_2` must have dimensions + "inversion_sigma", "spherical_harmonics_l", "species_center", and any "l{x}" + keys (where x is a positive integer) leftover from previous CG tensor + products. + + This function is an iterative special case of the more general + :py:mod:`correlate_tensors`. As a density is being correlated with itself, + some redundant CG tensor products can be skipped with the `skip_redundant` + keyword. + + :param tensor_1: the first tensor to correlate. + :param tensor_2: the second tensor to correlate. + :param selected_keys: :py:class:`Labels` or `List[:py:class:`Labels`]` + specifying the angular and/or parity channels to output at each + iteration. All :py:class:`Labels` objects passed here must only contain + key names "spherical_harmonics_l" and "inversion_sigma". If a single + :py:class:`Labels` object is passed, this is applied to the final + iteration only. If a :py:class:`list` of :py:class:`Labels` objects is + passed, each is applied to its corresponding iteration. If None is + passed, all angular and parity channels are output at each iteration, + with the global `angular_cutoff` applied if specified. + + :return: A :py:class:`TensorMap` corresponding to the correlated tensors. + """ + return _correlate_tensors( + tensor_1, + tensor_2, + selected_keys, + compute_metadata_only=False, + sparse=True, # sparse CG cache by default + ) + + +def correlate_tensors_metadata( + tensor_1: TensorMap, + tensor_2: TensorMap, + selected_keys: Optional[Union[Labels, List[Labels]]] = None, +) -> TensorMap: + """ + Returns the metadata-only :py:class:`TensorMap`(s) that would be output by + the function :py:func:`correlate_tensors` under the same settings, without + perfoming the actual Clebsch-Gordan tensor product. See this function for + full documentation. + """ + return _correlate_tensors( + tensor_1, + tensor_2, + selected_keys, + compute_metadata_only=True, + sparse=True, # sparse CG cache by default + ) + + +# ==================================================================== +# ===== Private functions that do the work on the TensorMap level +# ==================================================================== + + +def _correlate_tensors( + tensor_1: TensorMap, + tensor_2: TensorMap, + selected_keys: Optional[Union[Labels, List[Labels]]] = None, + compute_metadata_only: bool = False, + sparse: bool = True, +) -> TensorMap: + """ + Performs the Clebsch-Gordan tensor product for public functions + :py:func:`correlate_tensors` and :py:func:`correlate_tensors_metadata`. + """ + # Check inputs + # TODO: implement combinations of gradients too + for tensor in [tensor_1, tensor_2]: + if _dispatch.any([len(list(block.gradients())) > 0 for block in tensor]): + raise NotImplementedError( + "Clebsch Gordan combinations with gradients not yet implemented." + " Use metatensor.remove_gradients to remove gradients from the input." + ) + # # Check metadata + # if not ( + # _dispatch.all(tensor.keys.names[:2] == ["spherical_harmonics_l", "species_center"]) + # or _dispatch.all( + # tensor.keys.names[:3] + # == ["spherical_harmonics_l", "species_center", "species_neighbor"] + # ) + # ): + # raise ValueError( + # "input tensors must have key names" + # ' ["spherical_harmonics_l", "species_center"] or' + # ' ["spherical_harmonics_l", "species_center", "species_neighbor"]' + # ' as the first two or three keys' + # ) + if not _dispatch.all(tensor.component_names == ["spherical_harmonics_m"]): + raise ValueError( + "input tensors must have a single component" + " axis with name `spherical_harmonics_m`" + ) + # tensor_1 = _clebsch_gordan._standardize_keys(tensor_1) + # tensor_2 = _clebsch_gordan._standardize_keys(tensor_2) + + # Parse the selected keys + selected_keys = _clebsch_gordan._parse_selected_keys( + n_iterations=1, + selected_keys=selected_keys, + like=tensor_1.keys.values, + ) + + # Pre-compute the keys needed to perform each CG tensor product + key_metadata = _clebsch_gordan._precompute_keys( + tensor_1.keys, + tensor_2.keys, + n_iterations=1, + selected_keys=selected_keys, + skip_redundant=[False], + )[0] + + # Compute CG coefficient cache + if compute_metadata_only: + cg_cache = None + else: + angular_max = max( + _dispatch.concatenate( + [tensor_1.keys.column("spherical_harmonics_l")] + + [tensor_2.keys.column("spherical_harmonics_l")] + + [key_metadata[2].column("spherical_harmonics_l")] + ) + ) + # TODO: keys have been precomputed, so perhaps we don't need to + # compute all CG coefficients up to angular_max here. + # TODO: use sparse cache by default until we understand under which + # circumstances (and if) dense is faster. + cg_cache = _cg_cache.ClebschGordanReal(angular_max, sparse=sparse) + + # Perform CG tensor product by combining block pairs + blocks_out = [] + for key_1, key_2, key_out in zip(*key_metadata): + block_out = _clebsch_gordan._combine_blocks_same_samples( + tensor_1[key_1], + tensor_2[key_2], + key_out["spherical_harmonics_l"], + cg_cache, + compute_metadata_only=compute_metadata_only, + ) + blocks_out.append(block_out) + + # Build the TensorMap + return TensorMap(keys=key_metadata[2], blocks=blocks_out) diff --git a/python/rascaline/tests/utils/clebsch_gordan.py b/python/rascaline/tests/utils/correlate_density.py similarity index 90% rename from python/rascaline/tests/utils/clebsch_gordan.py rename to python/rascaline/tests/utils/correlate_density.py index 38b7fab32..edbcca160 100644 --- a/python/rascaline/tests/utils/clebsch_gordan.py +++ b/python/rascaline/tests/utils/correlate_density.py @@ -10,9 +10,9 @@ import rascaline from rascaline.utils import PowerSpectrum from rascaline.utils.clebsch_gordan._cg_cache import ClebschGordanReal -from rascaline.utils.clebsch_gordan.clebsch_gordan import ( +from rascaline.utils.clebsch_gordan._clebsch_gordan import _standardize_keys +from rascaline.utils.clebsch_gordan.correlate_density import ( _correlate_density, - _standardize_keys, correlate_density, correlate_density_metadata, ) @@ -43,9 +43,9 @@ DATA_ROOT = os.path.join(os.path.dirname(__file__), "data") SPHEX_HYPERS = { - "cutoff": 3.0, # Angstrom - "max_radial": 6, # Exclusive - "max_angular": 4, # Inclusive + "cutoff": 2.5, # Angstrom + "max_radial": 3, # Exclusive + "max_angular": 3, # Inclusive "atomic_gaussian_width": 0.2, "radial_basis": {"Gto": {}}, "cutoff_function": {"ShiftedCosine": {"width": 0.5}}, @@ -53,7 +53,7 @@ } SPHEX_HYPERS_SMALL = { - "cutoff": 3.0, # Angstrom + "cutoff": 2.5, # Angstrom "max_radial": 1, # Exclusive "max_angular": 2, # Inclusive "atomic_gaussian_width": 0.2, @@ -160,19 +160,7 @@ def get_norm(tensor: TensorMap): ) @pytest.mark.parametrize( "frames, nu_target, angular_cutoff, selected_keys", - [ - ( - h2_isolated(), - 3, - None, - Labels( - names=["spherical_harmonics_l", "inversion_sigma"], - values=np.array([[0, 1], [4, 1], [5, 1]]), - ), - ), - (h2o_isolated(), 2, 5, None), - (h2o_periodic(), 2, 5, None), - ], + [(h2o_periodic(), 2, 3, None)], ) def test_so3_equivariance(frames, nu_target, angular_cutoff, selected_keys): """ @@ -208,19 +196,7 @@ def test_so3_equivariance(frames, nu_target, angular_cutoff, selected_keys): ) @pytest.mark.parametrize( "frames, nu_target, angular_cutoff, selected_keys", - [ - ( - h2_isolated(), - 3, - None, - Labels( - names=["spherical_harmonics_l"], - values=np.array([0, 4, 5]).reshape(-1, 1), - ), - ), - (h2o_isolated(), 2, 5, None), - (h2o_periodic(), 2, 5, None), - ], + [(h2_isolated(), 2, 3, None)], ) def test_o3_equivariance(frames, nu_target, angular_cutoff, selected_keys): """ @@ -259,10 +235,7 @@ def test_o3_equivariance(frames, nu_target, angular_cutoff, selected_keys): @pytest.mark.parametrize("frames", [h2_isolated()]) @pytest.mark.parametrize( "sphex_powspec", - [ - (spherical_expansion, power_spectrum), - (spherical_expansion_small, power_spectrum_small), - ], + [(spherical_expansion, power_spectrum)], ) def test_lambda_soap_vs_powerspectrum(frames, sphex_powspec): """ @@ -282,7 +255,10 @@ def test_lambda_soap_vs_powerspectrum(frames, sphex_powspec): names=["spherical_harmonics_l"], values=np.array([0]).reshape(-1, 1) ), ) - keys = lsoap.keys.remove(name="spherical_harmonics_l") + keys = lsoap.keys + keys = keys.remove(name="order_nu") + keys = keys.remove(name="inversion_sigma") + keys = keys.remove(name="spherical_harmonics_l") lsoap = TensorMap(keys=keys, blocks=[b.copy() for b in lsoap.blocks()]) # Manipulate metadata to match that of PowerSpectrum: @@ -303,11 +279,6 @@ def test_lambda_soap_vs_powerspectrum(frames, sphex_powspec): ) ) lsoap = TensorMap(keys=lsoap.keys, blocks=blocks) - - # Compare metadata - assert metatensor.equal_metadata(lsoap, ps) - - # allclose on values assert metatensor.allclose(lsoap, ps) @@ -317,8 +288,8 @@ def test_lambda_soap_vs_powerspectrum(frames, sphex_powspec): @pytest.mark.skipif( not HAS_METATENSOR_OPERATIONS, reason="metatensor-operations is not installed" ) -@pytest.mark.parametrize("frames", [h2_isolated(), h2o_periodic()]) -@pytest.mark.parametrize("correlation_order", [2, 3, 4]) +@pytest.mark.parametrize("frames", [h2o_periodic()]) +@pytest.mark.parametrize("correlation_order", [2, 4]) def test_correlate_density_norm(frames, correlation_order): """ Checks \\|ρ^\\nu\\| = \\|ρ\\|^\\nu in the case where l lists are not @@ -434,7 +405,7 @@ def test_clebsch_gordan_orthogonality(cg_cache_dense, l1, l2): @pytest.mark.skipif( not HAS_METATENSOR_OPERATIONS, reason="metatensor-operations is not installed" ) -@pytest.mark.parametrize("frames", [h2_isolated(), h2o_isolated()]) +@pytest.mark.parametrize("frames", [h2o_periodic()]) def test_correlate_density_dense_sparse_agree(frames): """ Tests for agreement between nu=3 tensors built using both sparse and dense @@ -446,13 +417,13 @@ def test_correlate_density_dense_sparse_agree(frames): # sparse v dense CG cache n_body_sparse = _correlate_density( density, - correlation_order=3, + correlation_order=2, compute_metadata_only=False, sparse=True, ) n_body_dense = _correlate_density( density, - correlation_order=3, + correlation_order=2, compute_metadata_only=False, sparse=False, ) @@ -467,8 +438,8 @@ def test_correlate_density_dense_sparse_agree(frames): not HAS_METATENSOR_OPERATIONS, reason="metatensor-operations is not installed" ) @pytest.mark.parametrize("frames", [h2o_isolated()]) -@pytest.mark.parametrize("correlation_order", [2, 3]) -@pytest.mark.parametrize("skip_redundant", [True, False]) +@pytest.mark.parametrize("correlation_order", [3]) +@pytest.mark.parametrize("skip_redundant", [True]) def test_correlate_density_metadata_agree(frames, correlation_order, skip_redundant): """ Tests that the metadata of outputs from :py:func:`correlate_density` and @@ -479,7 +450,7 @@ def test_correlate_density_metadata_agree(frames, correlation_order, skip_redund nux = correlate_density( nu1, correlation_order=correlation_order, - angular_cutoff=None, + angular_cutoff=3, selected_keys=None, skip_redundant=skip_redundant, ) @@ -488,7 +459,7 @@ def test_correlate_density_metadata_agree(frames, correlation_order, skip_redund nux_metadata_only = correlate_density_metadata( nu1, correlation_order=correlation_order, - angular_cutoff=None, + angular_cutoff=3, selected_keys=None, skip_redundant=skip_redundant, ) @@ -501,7 +472,7 @@ def test_correlate_density_metadata_agree(frames, correlation_order, skip_redund [ None, Labels( - names=["spherical_harmonics_l"], values=np.array([1, 2, 4]).reshape(-1, 1) + names=["spherical_harmonics_l"], values=np.array([1, 3]).reshape(-1, 1) ), ], ) diff --git a/python/rascaline/tests/utils/correlate_tensors.py b/python/rascaline/tests/utils/correlate_tensors.py new file mode 100644 index 000000000..ddc3ea575 --- /dev/null +++ b/python/rascaline/tests/utils/correlate_tensors.py @@ -0,0 +1,176 @@ +# -*- coding: utf-8 -*- +import os +from typing import List + +import metatensor +import numpy as np +import pytest +from metatensor import Labels, TensorBlock, TensorMap + +import rascaline +from rascaline.utils import PowerSpectrum +from rascaline.utils.clebsch_gordan._cg_cache import ClebschGordanReal +from rascaline.utils.clebsch_gordan._clebsch_gordan import _standardize_keys +from rascaline.utils.clebsch_gordan.correlate_tensors import ( + _correlate_tensors, + correlate_tensors, + correlate_tensors_metadata, +) +from rascaline.utils.clebsch_gordan.correlate_density import correlate_density + + +# Try to import some modules +ase = pytest.importorskip("ase") +import ase.io # noqa: E402 + + +try: + import metatensor.operations + + HAS_METATENSOR_OPERATIONS = True +except ImportError: + HAS_METATENSOR_OPERATIONS = False +try: + import sympy # noqa: F401 + + HAS_SYMPY = True +except ImportError: + HAS_SYMPY = False + +if HAS_SYMPY: + from .rotations import WignerDReal, transform_frame_o3, transform_frame_so3 + + +DATA_ROOT = os.path.join(os.path.dirname(__file__), "data") + +SPHEX_HYPERS = { + "cutoff": 3.0, # Angstrom + "max_radial": 6, # Exclusive + "max_angular": 4, # Inclusive + "atomic_gaussian_width": 0.2, + "radial_basis": {"Gto": {}}, + "cutoff_function": {"ShiftedCosine": {"width": 0.5}}, + "center_atom_weight": 1.0, +} + +SPHEX_HYPERS_SMALL = { + "cutoff": 3.0, # Angstrom + "max_radial": 1, # Exclusive + "max_angular": 2, # Inclusive + "atomic_gaussian_width": 0.2, + "radial_basis": {"Gto": {}}, + "cutoff_function": {"ShiftedCosine": {"width": 0.5}}, + "center_atom_weight": 1.0, +} + + +# ============ Pytest fixtures ============ + + +@pytest.fixture() +def cg_cache_sparse(): + return ClebschGordanReal(lambda_max=5, sparse=True) + + +@pytest.fixture() +def cg_cache_dense(): + return ClebschGordanReal(lambda_max=5, sparse=False) + + +# ============ Helper functions ============ + + +# def h2_isolated(): +# return ase.io.read(os.path.join(DATA_ROOT, "h2_isolated.xyz"), ":") + + +def h2o_isolated(): + return ase.io.read(os.path.join(DATA_ROOT, "h2o_isolated.xyz"), ":") + + +# def h2o_periodic(): +# return ase.io.read(os.path.join(DATA_ROOT, "h2o_periodic.xyz"), ":") + + +# def wigner_d_matrices(lmax: int): +# return WignerDReal(lmax=lmax) + + +# def spherical_expansion(frames: List[ase.Atoms]): +# """Returns a rascaline SphericalExpansion""" +# calculator = rascaline.SphericalExpansion(**SPHEX_HYPERS) +# return calculator.compute(frames) + + +def spherical_expansion_small(frames: List[ase.Atoms]): + """Returns a rascaline SphericalExpansion""" + calculator = rascaline.SphericalExpansion(**SPHEX_HYPERS_SMALL) + return calculator.compute(frames) + + +# def power_spectrum(frames: List[ase.Atoms]): +# """Returns a rascaline PowerSpectrum constructed from a +# SphericalExpansion""" +# return PowerSpectrum(rascaline.SphericalExpansion(**SPHEX_HYPERS)).compute(frames) + + +# def power_spectrum_small(frames: List[ase.Atoms]): +# """Returns a rascaline PowerSpectrum constructed from a +# SphericalExpansion""" +# return PowerSpectrum(rascaline.SphericalExpansion(**SPHEX_HYPERS_SMALL)).compute( +# frames +# ) + + +# def get_norm(tensor: TensorMap): +# """ +# Calculates the norm used in CG iteration tests. Assumes that the TensorMap +# is sliced to a single sample. + +# For a given atomic sample, the norm is calculated for each feature vector, +# as a sum over lambda, sigma, and m. +# """ +# # Check that there is only one sample +# assert ( +# len( +# metatensor.unique_metadata( +# tensor, "samples", ["structure", "center", "species_center"] +# ).values +# ) +# == 1 +# ) +# norm = 0.0 +# for key, block in tensor.items(): # Sum over lambda and sigma +# angular_l = key["spherical_harmonics_l"] +# norm += np.sum( +# [ +# np.linalg.norm(block.values[0, m, :]) ** 2 +# for m in range(-angular_l, angular_l + 1) +# ] +# ) + +# return norm + + +# ============ Test equivalence with `correlate_density` ============ + +@pytest.mark.skipif( + not HAS_METATENSOR_OPERATIONS, reason="metatensor-operations is not installed" +) +@pytest.mark.parametrize("frames", [h2o_isolated()]) +@pytest.mark.parametrize("selected_keys", [None]) +def test_correlate_tensors_correlate_density_equivalent(frames, selected_keys): + """ + Tests that using `correlate_tensors` with two identical densities + """ + density = spherical_expansion_small(frames) + + # NOTE: testing the private function here so we can control the use of + # sparse v dense CG cache + correlated_density = correlate_density(density, correlation_order=2, selected_keys=selected_keys) + correlated_tensor = correlate_tensors( + density, density, selected_keys=selected_keys + ) + correlated_tensor = correlated_tensor.keys_to_properties(keys_to_move=["l1", "l2"]) + + assert metatensor.equal(correlated_density, correlated_tensor)