Dtype issue with lambda-SOAP

[frostedoyster]

Attaching the script here

import featomic.torch
from featomic.torch import SphericalExpansion, systems_to_torch
from featomic.torch.clebsch_gordan import EquivariantPowerSpectrum
import ase.io
import torch
torch.set_default_dtype(torch.float64)


HYPERPARAMETERS = {
    "cutoff": {
        "radius": 5.0,
        "smoothing": {"type": "ShiftedCosine", "width": 0.5},
    },
    "density": {
        "type": "Gaussian",
        "width": 0.3,
    },
    "basis": {
        "type": "TensorProduct",
        "max_angular": 6,
        "radial": {"type": "Gto", "max_radial": 4},
    },
}


spex_calculator = SphericalExpansion(**HYPERPARAMETERS)
calculator = EquivariantPowerSpectrum(spex_calculator)


systems = ase.io.read("qm9_reduced_100.xyz", ":")
systems = systems_to_torch(systems)

power_spectrum = calculator.compute(systems, neighbors_to_properties=True)

This doesn't work without the torch.set_default_dtype(torch.float64) line. In that case, one gets

Traceback (most recent call last):
  File "generate_lambda_soap/dump.py", line 31, in <module>
    power_spectrum = calculator.compute(systems, neighbors_to_properties=True)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "lib/python3.12/site-packages/featomic/clebsch_gordan/_equivariant_power_spectrum.py", line 260, in compute
    return self._equivariant_power_spectrum(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "lib/python3.12/site-packages/featomic/clebsch_gordan/_equivariant_power_spectrum.py", line 370, in _equivariant_power_spectrum
    pow_spec = self._cg_product.compute(
               ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "lib/python3.12/site-packages/featomic/clebsch_gordan/_cg_product.py", line 204, in compute
    return self._cg_tensor_product(
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "lib/python3.12/site-packages/featomic/clebsch_gordan/_cg_product.py", line 323, in _cg_tensor_product
    _utils.cg_tensor_product_blocks(
  File "lib/python3.12/site-packages/featomic/clebsch_gordan/_utils.py", line 245, in cg_tensor_product_blocks
    combined_values = _coefficients.cg_tensor_product(
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/filippo/code/virtualenvs/base/lib/python3.12/site-packages/featomic/clebsch_gordan/_coefficients.py", line 650, in cg_tensor_product
    return _cg_tensor_product_dense(array_1, array_2, o3_lambdas, cg_coefficients)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "lib/python3.12/site-packages/featomic/clebsch_gordan/_coefficients.py", line 748, in _cg_tensor_product_dense
    output = _cg_couple_dense(tensor_product, o3_lambda, cg_coefficients)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "lib/python3.12/site-packages/featomic/clebsch_gordan/_coefficients.py", line 582, in _cg_couple_dense
    return _dispatch.tensordot(array, cg_l1l2lam[0, ..., 0], axes=([2, 1], [1, 0]))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "lib/python3.12/site-packages/featomic/clebsch_gordan/_dispatch.py", line 87, in tensordot
    return torch.tensordot(array_1, array_2, axes)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "lib/python3.12/site-packages/torch/functional.py", line 1382, in tensordot
    return _VF.tensordot(a, b, dims_a, dims_b)  # type: ignore[attr-defined]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: both inputs should have same dtype

[Luthaf]

Ok, the easy fix here is to pass a dtype to EquivariantPowerSpectrum, i.e. calculator = EquivariantPowerSpectrum(spex_calculator, dtype=torch.float64)

It would also be good to have a way to control the dtype that comes out of featomic.torch.systems_to_torch, and add a dtype argument there as well. The function from metatomic handles this, so you can use it instead (https://docs.metatensor.org/metatomic/latest/torch/reference/systems.html#metatomic.torch.systems_to_torch).