Confusing "parameter flow" for the calculators

[ceriottm]

The calculators use a mix of class parameters and "compute function" parameters to define e.g. the smearing and the calculation accuracy. This is confusing and a bit clunky - it'd be better to streamline the implementation to avoid having to worry too much about (re)setting these values.

[sirmarcel]

Just to add to this, it's a huge potential foot-gun that various parts of the infrastructure now store the cell/the mesh and need to be kept updated. It may be worth thinking about converting everything into something functionally pure and then implementing caching based on these guarantees -- but I don't know yet how to best do this. I agree that this is a problem that needs to be addressed.

[PicoCentauri]

Not directly the issue but related:

While benchmarking, I found a couple of places that are slowing down the code just because we do multiplication with floats in the code. @sirmarcel already mentioned this but now I really saw it.

On top: In many places we may overcome repeating calculations by adding buffers in the init of each class to compute prefactors.

[ceriottm]

On top: In many places we may overcome repeating calculations by adding buffers in the init of each class to compute prefactors.

We have to be careful here if the prefactors depend on stuff, e.g. the smearing.

[PicoCentauri]

Good point!

this is why these points are in the correct place inside this issue. We can also update buffers if we have to, but hopefully we can avoid.

[PicoCentauri]

We had some brainstorming meeting with @E-Rum and @GardevoirX and came up some refined structure for the calculator classes.

Basically, we will split the calculators into classes of EwaldCalculator (Ewald, PME, P3M) and a single DirectCalculator.

This means we don't have to branch a lot if a potential has a finite smearing or not. There will only be a single check at the init:

import torch


class CalcultorBase(torch.nn.Module):
    def __init__(self): ...

    def forward(self):
        # validate the input and call the _compute method
        return self._compute()

    def _compute(self):
        raise NotImplementedError("Only implemented in subclasses")


class EwaldCalculator(CalcultorBase):
    def __init__(self, potential):
        super().__init__()

        if potential.smearing is None:
            raise ValueError()

    def _compute(self):
        return (
            self._compute_kspace()
            + self._compute_rspace()
            + self._compute_background_contribution()
            + self._compute_self_contriubution()
        )

    def _compute_background_contribution(self): ...

    def _compute_self_contriubution(self): ...

    def _compute_kspace(self): ...

    def _compute_rspace(self): ...


class PMECalculator(EwaldCalculator):
    def __init__(self): ...

    def _compute_kspace(self):
        # override `_compute_kspace` from Ewald with PME logic
        ...


class P3MCalculator(EwaldCalculator):
    def __init__(self): ...

    def _compute_kspace(self): ...


class DirectCalculator(CalcultorBase):
    def __init__(self, potential):

        if potential.smearing is not None:
            raise ValueError()

    def _compute(self):
        # call potential.from_dist() and multiply the charges
        ...

As you see this will not break API at all and should simplify the parameter flow. If there are no objection, @E-Rum volunteered to check the implementation and do it. THANKS A LOT ALREADY!!!!