subclassing ndarray

lyncs_quda/clover_field.py

@@ -6,7 +6,7 @@
     "CloverField",
 ]
 
-import numpy
+import numpy, inspect
 from cppyy.gbl.std import vector
 
 from lyncs_cppyy import make_shared, to_pointer
@@ -34,89 +34,109 @@ class CloverField(LatticeField):
      * Only rho is mutable.  To change other params, a new instance should be created
      * QUDA convention for clover field := 1+i ( kappa csw )/4 sigma_mu,nu F_mu,nu (<-sigma_mu,nu: spinor tensor)
      *  so that sigma_mu,nu = i[g_mu, g_nu], F_mu,nu = (Q_mu,nu - Q_nu,mu)/8 (1/2 is missing from sigma_mu,nu)
+     * Apparently, an input to QUDA clover object, coeff = kappa*csw 
+     *  wihout a normalization factor of 1/4 or 1/32 (suggested in interface_quda.cpp) 
     """
 
+    _children = {}
+
+    def __new__(cls, fmunu, **kwargs):
+        #TODO: get dofs and local dims from kwargs, instead of getting them
+        # from self.shape assuming that it has the form (dofs, local_dims)
+        if isinstance(fmunu, CloverField):
+            return fmunu
+        if not isinstance(fmunu, (numpy.ndarray, cupy.ndarray)):
+            raise TypeError(
+                f"Supporting only numpy or cupy for field, got {type(fmunu)}"
+            )
+        is_clover = False
+        if not isinstance(fmunu, GaugeField):
+            if kwargs.get("is_clover", False):
+                is_clover = True
+                parent = type(fmunu)
+                field = fmunu
+            else:
+                fmunu = GaugeField(fmunu)
+
+        # Set parent
+        if not is_clover: # not copying from a clover-field array
+            idof = int((fmunu.ncol * fmunu.ndims) ** 2 / 2)
+            prec = fmunu.dtype
+            field = fmunu.backend.empty((idof,) + fmunu.dims, dtype=prec)
+            parent = fmunu.backend.ndarray
+        # Set child
+        if parent in cls._children.keys():
+            child  = cls._children.get(parent)
+        else:
+            child = type(cls.__name__+"ext",(cls, parent), {})
+            cls._children.update({parent: child})
+        obj = field.view(type=child)
+
+        #self._dims = kwargs.get("dims", self.shape[-self.ndims :])
+	#self._dofs = kwargs.get("dofs", field.shape[: -self.ndims])
+
+        return obj
+
     def __init__(
         self,
-        fmunu,
+        obj,
         coeff=0.0,
         twisted=False,
         mu2=0,
         tf="NO",
         eps2=0,
         rho=0,
         computeTrLog=False,
+        **kwargs
     ):
-        # ASSUME: coeff = kappa*csw wihout a normalization factor of 1/4 or 1/32 (suggested in interface_quda.cpp)
-        # ? better to store fmunu.quda_field to _fmunu -> import gauge_tensor to be used in some methods
-        # ? better to put clover into self.field -> need walk-around to make copy() work
-        if not isinstance(fmunu, GaugeField):
-            fmunu = GaugeField(fmunu)
-        self._fmunu = fmunu.compute_fmunu()
-        super().__init__(self._fmunu.field, comm=self._fmunu.comm)
-
-        # QUDA clover field inherently works with real's not with complex's (c.f., include/clover_field_order.h)
-
-        idof = int((self._fmunu.ncol * self._fmunu.ndims) ** 2 / 2)
-        prec = self._fmunu.precision
-        self._direct = (
-            False  # Here, it is a flag to indicate whether the field has been computed
-        )
-        self._inverse = (
-            False  # Here, it is a flag to indicate whether the field has been computed
-        )
-
-        new = lambda idof: LatticeField.create(
-            self._fmunu.global_lattice,
-            dofs=(idof,),
-            dtype=prec,
-            device=self._fmunu.device,
-            empty=True,
-        )
-        self._clover = new(idof)
-        self._cloverInv = new(idof)
-        self.coeff = coeff
-        self._twisted = twisted
-        self._twist_flavor = tf
-        self._mu2 = mu2
-        self._eps2 = eps2
-        self._rho = rho
-        self.computeTrLog = computeTrLog
-
-    # shape, dofs, dtype, iscomlex, isreal are overwriten to report their values for the clover field, instead of _fmunu
-
-    @property
-    def shape(self):
-        "Shape of the clover field"
-        return self._clover.shape
-
-    @property
-    def dofs(self):
-        "Shape of the per-site degrees of freedom"
-        return self._clover.dofs
-
-    @property
-    def dtype(self):
-        "Clover field data type"
-        return self._clover.dtype
-
-    @property
-    def iscomplex(self):
-        "Whether the clover field dtype is complex"
-        return self._clover.iscomplex
-
-    @property
-    def isreal(self):
-        "Whether the clover field dtype is real"
-        return self._clover.isreal
-
+        # WARNING: ndarray object is not supposed to be view-casted to CloverField object
+        #           except in __new__, for which __init__ will be called subsequently,
+        #           as the result won't come with  meta info such as 'coeff' or 'mu2'
+
+        super().__init__(obj, getattr(obj, "comm", None))
+        if isinstance(obj, GaugeField):
+            # explicit construction
+            # QUDA clover field inherently works with real's not with complex's (c.f., include/clover_field_order.h)
+            self._cloverInv = LatticeField.create(
+                self.global_lattice,
+                dofs=self.dofs,
+                dtype=self.dtype,
+                device=self.device_id,
+                empty=True,
+            )
+            self._fmunu = obj.compute_fmunu()
+            self._direct = (
+                False  # Here, it is a flag to indicate whether the field has been computed
+            )
+            self._inverse = (
+                False  # Here, it is a flag to indicate whether the field has been computed
+            )
+            self.coeff = coeff
+            self._twisted = twisted
+            self._twist_flavor = tf
+            self._mu2 = mu2
+            self._eps2 = eps2
+            self._rho = rho
+            self.computeTrLog = computeTrLog
+        elif isinstance(obj, CloverField):
+            # upcasting to ndarray or new from template
+            self.__dict__.update(obj.__dict__)
+        elif isinstance(obj, self.backend.ndarray):
+            pass
+        else:
+            raise ValueError("The input is expected to be ndarray or LatticeField object")
+
+    def _prepare(self, field, copy=False, check=False, **kwargs):
+        # When CloverField object prepares its input, the input is assumed to be of CloverField
+        return super()._prepare(field, copy=copy, check=check, is_clover=True, **kwargs)
+
     # naming suggestion: native_view? default_* saved for dofs+lattice?
     def default_view(self):
         N = 1 if self.order == "FLOAT2" else 4
         shape = (2,)  # even-odd
         shape += (self.dofs[0] // N, -1, N)
 
-        return self.field.view().reshape(shape)
+        return self.float_view().reshape(shape)
 
     @property
     def twisted(self):
@@ -178,7 +198,7 @@ def quda_params(self):
         params = lib.CloverFieldParam()
         lib.copy_struct(params, super().quda_params)
         params.inverse = True
-        params.clover = to_pointer(self._clover.ptr)
+        params.clover = to_pointer(self.ptr)
         params.cloverInv = to_pointer(self._cloverInv.ptr)
         params.coeff = self.coeff
         params.twisted = self.twisted
@@ -205,15 +225,15 @@ def clover_field(self):
         if not self._direct:
             lib.computeClover(self.quda_field, self._fmunu.quda_field, self.coeff)
             self._direct = True
-        return self._clover.field
+        return self.view(self.backend.ndarray)
 
     @property
     def inverse_field(self):
         if not self._inverse:
             self.clover_field
             lib.cloverInvert(self.quda_field, self.computeTrLog)
             self._inverse = True
-        return self._cloverInv.field
+        return self._cloverInv.view(self.backend.ndarray)
 
     @property
     def trLog(self):

lyncs_quda/dirac.py

@@ -242,9 +242,9 @@ def action(self, phi, **params):
         out = 0
         if not self.full and "CLOVER" in self.type:
             if self.even:
-                out -= 2 * self.clover.trLog[1]
-            else:
                 out -= 2 * self.clover.trLog[0]
+            else:
+                out -= 2 * self.clover.trLog[1]
 
         parity = None
         if not self.full:

lyncs_quda/gauge_field.py

@@ -16,7 +16,7 @@
 from time import time
 from math import sqrt
 from collections import defaultdict
-import numpy
+import numpy, inspect
 from lyncs_cppyy import make_shared, lib as tmp, to_pointer, array_to_pointers
 from lyncs_utils import prod, isiterable
 from .lib import lib, cupy
@@ -65,11 +65,21 @@ def momentum(lattice, **kwargs):
 class GaugeField(LatticeField):
     "Mimics the quda::GaugeField object"
 
-    @LatticeField.field.setter
-    def field(self, field):
-        LatticeField.field.fset(self, field)
-        if self.reconstruct == "INVALID":
-            raise TypeError(f"Unrecognized field dofs {self.dofs}")
+    _children = {}
+
+    def _check_field(self, field=None):
+        # NOTE:
+        #  - For now, we store dofs and local lattice info in the shape of the array
+        #  - The canonical form for GaugeField is [dofs, local_lattice]
+        #     where len(dofs) == 1,2 and len(local_lattice) = 4
+        super()._check_field(field)
+        if field is None:
+            field = self
+        dofs = field.shape[:-self.ndims]
+        pdofs = prod(dofs) if dofs[0] not in self._geometry_values[1] or dofs[0]==1 else prod(dofs[1:])
+        pdofs *= 2**(self.iscomplex)
+        if not (pdofs in (12, 8, 10) or sqrt(pdofs / 2).is_integer()):
+            raise TypeError(f"Unrecognized field dofs {dofs}")
 
     def new(self, reconstruct=None, geometry=None, **kwargs):
         "Returns a new empty field based on the current"
@@ -130,11 +140,6 @@ def copy(self, other=None, out=None, **kwargs):
 
         return super().copy(other, out, **kwargs)
 
-    def __array_finalize__(self, obj):
-        "Support for __array_finalize__ standard"
-        # need to reset QUDA object when meta data of its Python wrapper is changed
-        self._quda = None
-
     def _prepare(self, field, **kwargs):
         field = super()._prepare(field, **kwargs)
         is_momentum = kwargs.get("is_momentum", self.is_momentum)
@@ -245,7 +250,7 @@ def is_momentum(self):
     def is_momentum(self, value):
         self._is_momentum = value
         if self._quda is not None:
-            self._quda.link_type = self.quda_link_type
+            self._quda = None
 
     @property
     def t_boundary(self):
@@ -501,7 +506,7 @@ def project(self, tol=None):
         if tol is None:
             tol = numpy.finfo(self.dtype).eps
 
-        assert self.device == cupy.cuda.runtime.getDevice()
+        assert self.device_id == cupy.cuda.runtime.getDevice()
         fails = cupy.zeros((1,), dtype="int32")
         lib.projectSU3(self.quda_field, tol, to_pointer(fails.data.ptr, "int *"))
         # return fails.get()[0]  # shouldn't we reduce?