Add Radial and initial Grid

Author: oerc0122

MDANSE/Src/MDANSE/Core/SubclassFactory.py

@@ -13,12 +13,14 @@
 #    You should have received a copy of the GNU General Public License
 #    along with this program.  If not, see <https://www.gnu.org/licenses/>.
 #
-
-from typing import TypeVar
 import difflib
-
+from abc import ABC
+from collections.abc import Callable, Sequence
+from typing import Dict, Generic, TypeVar, Union
 
 Self = TypeVar("Self", bound="SubclassFactory")
+T = TypeVar("T")
+
 # The Self TypeVar is a typing hint indicating that
 # a method of a class A will be returning an object
 # of type A as well. Since we don't know for which class
@@ -201,3 +203,75 @@ def indirect_subclass_dictionary(cls):
             dict(str:type) -- a dictionary of the subclasses of this class
         """
         return recursive_dict(cls)
+
+class RegisterFactory(ABC, Generic[T]):
+    """Alternative factory which uses explicit registration."""
+
+    _registered_subclasses: Dict[str, type[T]] = {}
+
+    @classmethod
+    def get(cls, name: str) -> type[T]:
+        """Get the class from the registry.
+
+        Parameters
+        ----------
+        name : str
+            Name of class to retrieve.
+
+        Returns
+        -------
+        type[T]
+            Class ready to instantiate.
+
+        Raises
+        ------
+        KeyError
+            Name not found in registry.
+        """
+        if name not in cls._registered_subclasses:
+            raise KeyError(f"Cannot instantiate class {name}, "
+                           f"not in registered subclasses ({','.join(cls._registered_subclasses)})")
+        return cls._registered_subclasses[name]
+
+    @classmethod
+    def create(cls, name: str, *args, **kwargs) -> T:
+        """Instantiate child class from registry.
+
+        Parameters
+        ----------
+        name : str
+            Name to instantiate.
+        """
+        return cls.get(name)(*args, **kwargs)
+
+    @classmethod
+    @property
+    def subclasses(cls):
+        return list(cls._registered_subclasses.keys())
+
+    @classmethod
+    def register(cls, names: Union[str, Sequence[str]]):
+        """Register class as factory member.
+
+        Parameters
+        ----------
+        names : str | Sequence[str]
+            Internal name(s) to use.
+
+        Examples
+        --------
+        .. code-block:: python
+            @ClassFactory.register("elem")
+            class MyClass(): ...
+
+        """
+        def class_wrapper(wrapped_class: type) -> Callable[..., T]:
+
+            if isinstance(names, str):
+                cls._registered_subclasses[names] = wrapped_class
+            else:
+                for name in names:
+                    cls._registered_subclasses[name] = wrapped_class
+            return wrapped_class
+
+        return class_wrapper

MDANSE/Src/MDANSE/Framework/NewQVectors/ArbitraryQVectors.py

@@ -0,0 +1,105 @@
+from typing import Optional
+
+import numpy as np
+from MDANSE.Framework.NewQVectors.QVector import (QVecGen,
+                                                  QVecGeneratorProtocol,
+                                                  QVectorData,
+                                                  QVectorGenerator)
+from MDANSE.MolecularDynamics.UnitCell import UnitCell
+from numpy.typing import ArrayLike
+
+@QVectorGenerator.register("ListQVectors")
+class ListQVectors(QVectorGenerator):
+    """Return Q vectors from a provided list of vectors.
+
+    Parameters
+    ----------
+    qvectors : ArrayLike
+        Sequence of vectors to return.
+    hkl : bool
+        Whether vectors are provided in reciprocal lattice units.
+    lattice : Optional[UnitCell]
+        Lattice to generate within.
+
+    Raises
+    ------
+    ValueError
+        If hkl and no lattice provided.
+
+    Examples
+    --------
+    >>> qvec = ListQVectors([[1, 2, 3], [2, 0, 5]])
+    >>> for i in qvec:
+    ...     print(i.q)
+    [1. 2. 3.]
+    [2. 0. 5.]
+    """
+
+    def __init__(self,
+                 qvectors: ArrayLike,
+                 *,
+                 lattice: Optional[UnitCell] = None,
+                 **kwargs):
+
+        super().__init__(lattice=lattice, **kwargs)
+        self.qvectors = np.array(qvectors)
+
+        if self.qvectors.shape != (len(self.qvectors), 3):
+            raise ValueError(f"`qvectors` ({qvectors.shape}) must be an (N, 3) sequence.")
+
+
+    def generate(self, lattice: Optional[UnitCell] = None) -> QVecGen:
+        lattice = lattice if lattice is not None else self.lattice
+        constructor = self.qvec_gen
+
+        while self._ind < len(self):
+            new_ind = yield constructor(self.qvectors[self._ind], lattice)
+
+            self._ind += 1
+            if new_ind is not None:
+                self.reset(new_ind)
+
+    def __len__(self) -> int:
+        return len(self.qvectors)
+
+    def reset(self, value: int = 0):
+        super().reset(value)
+
+
+@QVectorGenerator.register("GeneratorQVectors")
+class GeneratorQVectors(QVectorGenerator):
+    """Return Q Vectors as generated from a generator function.
+
+    Parameters
+    ----------
+    qvectors : QVecGen | Generator[ArrayLike, int, None]
+        Generator which returns QVectorData.
+    lattice : Optional[UnitCell]
+        Lattice in which to generate Q Vectors.
+
+    Examples
+    --------
+    FIXME: Add docs.
+
+    """
+    def __init__(self,
+                 qvectors: QVecGeneratorProtocol,
+                 *args,
+                 lattice: Optional[UnitCell] = None,
+                 returns_qvec_data: bool = True,
+                 **kwargs):
+        super().__init__(lattice=lattice, **kwargs)
+
+        self.generator = qvectors
+        self.args = args
+        self.returns_qvec_data = returns_qvec_data
+
+    def generate(self, lattice: Optional[UnitCell] = None) -> QVecGen:
+        lattice = lattice if lattice is not None else self.lattice
+        constructor = self.qvec_gen
+
+        for qvec in self.generator(*self.args, lattice):
+            yield qvec if self.returns_qvec_data else constructor(qvec, lattice)
+
+    def reset(self, val: None = None):
+        raise NotImplementedError("Cannot reset this generator")