MadSpin: generate decay events with mg7, and cut the overheads around them (2.0x at 10k, 4.4x at 100k)

MadSpin/decay.py

@@ -5427,6 +5427,9 @@ class DensityMatrix:
     # Cache diagonal masks by basis_id (depends only on helicities)
     _diag_cache = {}
 
+    # Same, as integer positions, for trace()
+    _diag_pos_cache = {}
+
     # Cache tensor-product helicity tables by basis_id
     _tp_hel_cache = {}
 
@@ -5629,6 +5632,21 @@ def _get_diag_mask_cached(self):
         DensityMatrix._diag_cache[self._basis_id] = mask
         return mask
 
+    def _get_diag_positions_cached(self):
+        """Positions of the diagonal entries, as a plain tuple of ints.
+
+        trace() sums a handful of entries -- two, for a single decaying
+        fermion -- and numpy costs about 0.9 us to do that however few there
+        are, nearly all of it dispatch. Indexing the cached positions directly
+        is 8x faster at that size. Cached per basis_id alongside the mask.
+        """
+        cached = DensityMatrix._diag_pos_cache.get(self._basis_id)
+        if cached is not None:
+            return cached
+        positions = tuple(int(i) for i in np.flatnonzero(self._diag_mask))
+        DensityMatrix._diag_pos_cache[self._basis_id] = positions
+        return positions
+
     # -------------------------------------------------------------------------
     # Cached permutation for alignment by helicity labels
     # -------------------------------------------------------------------------
@@ -5678,15 +5696,19 @@ def scalar_multiplication(self, other):
         # Fastest correct path for map-built matrices
         if (self.map_density_matrix_ind is not None and
                 self.map_density_matrix_ind is other.map_density_matrix_ind):
-            return np.sum(self.values * other.values)
+            # np.dot rather than np.sum(a*b): identical for complex (dot does
+            # not conjugate) but one call instead of two, and it skips the
+            # temporary the multiply would allocate. 0.23 us against 0.90 us
+            # on the 16-entry matrices this sees.
+            return np.dot(self.values, other.values)
 
         # Align by cached ordering for each basis
         self._ensure_sorted_view()
         other._ensure_sorted_view()
 
         a = self._sort_order
         b = other._sort_order
-        return np.sum(self.values[a] * other.values[b])
+        return np.dot(self.values[a], other.values[b])
 
     def tensor_product(self, other):
         """
@@ -5736,7 +5758,17 @@ def trace(self):
         """
         Order-independent trace.
         """
-        return np.sum(self.values[self._diag_mask])
+        # Sum the diagonal entries by position. See
+        # _get_diag_positions_cached: at these sizes numpy's dispatch dwarfs
+        # the addition itself.
+        values = self.values
+        positions = self._get_diag_positions_cached()
+        if not positions:
+            return np.complex64(0)
+        total = values[positions[0]]
+        for i in positions[1:]:
+            total = total + values[i]
+        return total
 
 
     def print_full_matrix(self, precision=6):