mod: use nested loops for resorting integrals

sympleints/Functions.py

@@ -27,7 +27,7 @@ class Functions:
     doc_func: Optional[Callable] = None
     comment: str = ""
     boys_func: Optional[str] = None
-    ncomponents: int = 0
+    ncomponents: int = 1
     with_ref_center: bool = True
     full_name: Optional["str"] = None
     l_aux_max: Optional[int] = None
@@ -39,7 +39,7 @@ class Functions:
     def __post_init__(self):
         assert self.l_max >= 0
         assert len(self.coeffs) == len(self.exponents) == len(self.centers)
-        assert self.ncomponents >= 0
+        assert self.ncomponents >= 1
 
         L_iter, inner = self.ls_exprs
         self.L_iter = L_iter

sympleints/PythonRenderer.py

@@ -27,7 +27,11 @@ class PythonRenderer(Renderer):
         "module": "py_module.tpl",
     }
     _primitive = False
-    _drop_dim = True
+    _drop_dim = True  # TODO: remove this?!
+    resort_func_dict = {
+        2: ("resort_ba_ab", (2, 1)),
+        3: ("resort_bac_abc", (2, 1)),
+    }
 
     def render_function(
         self,
@@ -84,27 +88,23 @@ def render_equi_function(
         equi_name,
         equi_inds,
         shape,
-        from_axes,
-        to_axes,
     ):
+        resort_func, new_order = self.resort_func_dict[functions.nbfs]
+        sizes = shape
+        sizes = [shape[i] for i in new_order]
         equi_args = ", ".join(functions.full_args_for_bf_inds(equi_inds))
         args = ", ".join(functions.full_args)
+        ncomponents = functions.ncomponents
+
         tpl = self.get_template(key="equi_func")
-        nbfs = functions.nbfs
-        assert nbfs in (2, 3), "Implement other cases in template!"
-        if functions.ncomponents == 1:
-            from_axes = tuple([i - 1 for i in from_axes[1:]])
-            to_axes = tuple([i - 1 for i in to_axes[1:]])
         rendered = tpl.render(
             equi_name=equi_name,
             equi_args=equi_args,
             name=name,
             args=args,
-            from_axes=from_axes,
-            to_axes=to_axes,
-            shape=shape,
-            functions=functions,
-            primitive=functions.primitive,
+            sizes=sizes,
+            resort_func=resort_func,
+            ncomponents=ncomponents,
         )
         return rendered
 
@@ -115,6 +115,7 @@ def render_func_dict(self, name, rendered_funcs):
 
     def render_module(self, functions, rendered_funcs, **tpl_kwargs):
         func_dict = self.render_func_dict(functions.name, rendered_funcs)
+        resort_func, _ = self.resort_func_dict[functions.nbfs]
         tpl = self.get_template(key="module")
         _tpl_kwargs = {
             "header": functions.header,
@@ -124,6 +125,7 @@ def render_module(self, functions, rendered_funcs, **tpl_kwargs):
             "func_dict": func_dict,
             "name": functions.name,
             "args": functions.full_args,
+            "resort_func": resort_func,
         }
         _tpl_kwargs.update(tpl_kwargs)
         rendered = tpl.render(**_tpl_kwargs)

sympleints/Renderer.py

@@ -74,6 +74,7 @@ def render_equi_function(
     def render_functions(self, functions: Functions):
         args = functions.full_args
         ncomponents = functions.ncomponents
+        assert ncomponents > 0
         if len(hermi_inds := functions.hermitian) >= 2:
             hermi_inds = functions.hermitian
             assert len(hermi_inds) == 2
@@ -124,34 +125,12 @@ def render_functions(self, functions: Functions):
                 L_tots_equi = tuple(L_tots_equi)
                 name_equi = func_name_from_Ls(functions.name, L_tots_equi)
 
-                from_axes = tuple(equi_inds)
-                to_axes = tuple(org_inds)
-
-                nbfs = functions.nbfs
-                # Axes/indices are missing
-                if len(from_axes) != nbfs:
-                    to_axes = tuple(range(nbfs))
-                    from_axes = list(range(nbfs))
-                    from_axes[hi], from_axes[hj] = from_axes[hj], from_axes[hi]
-                    from_axes = tuple(from_axes)
-
-                more_components = functions.ncomponents > 1
-                add_axis = more_components or (not self._drop_dim and nbfs == 2)
-                if add_axis:
-                    from_axes = tuple([0] + [fa + 1 for fa in from_axes])
-                    to_axes = tuple([0] + [ta + 1 for ta in to_axes])
-
-                # Drop first dimension when only 1 or 0 components are present
-                reshape = shape if more_components else shape[1:]
-
                 func_equi = self.render_equi_function(
                     functions,
                     name,
                     name_equi,
                     equi_inds,
-                    reshape,
-                    from_axes,
-                    to_axes,
+                    shape,
                 )
                 rendered_funcs.append(
                     RenderedFunction(name=name_equi, Ls=L_tots_equi, text=func_equi)

sympleints/templates/numba_equi_func.tpl

@@ -9,7 +9,6 @@ def {{ equi_name }}({{ args }}, result):
     """See docstring of {{ name }}."""
 
     # Call equivalent function and write to result
-    tmp = numpy.zeros_like(result)
-    {{ name }}({{ equi_args }}, tmp)
-    result[:]  {% if not primitive %}+{% endif %}= numpy.transpose(tmp.reshape({{ shape|join(", ") }}), axes={{ from_axes }}).flatten()
-
+    # tmp = numpy.zeros_like(result)
+    {{ name }}({{ equi_args }}, result)
+    {{ resort_func }}(result, {{ sizes|join(", ")}}, {{ ncomponents }})

sympleints/templates/numba_module.tpl

@@ -19,6 +19,29 @@ from pysisyphus.wavefunction.ints.boys import boys
 {{ ai }}
 {% endfor %}
 
+{% if resort_func == "resort_ba_ab" %}
+@numba.jit(
+    nopython=True,
+    nogil=True,
+    fastmath=True,
+    cache=True,
+)
+def resort_ba_ab(result, sizea, sizeb, ncomponents):
+    assert ncomponents > 0
+    tmp = numpy.zeros_like(result)
+    i = 0  # Original index
+    sizeab = sizea * sizeb
+    component_size = 0
+    for _ in range(ncomponents):
+        for b in range(sizeb):
+            for a in range(sizea):
+                j = component_size + (a * sizeb) + b  # New index
+                tmp[j] = result[i]
+                i += 1
+        component_size += sizeab
+    result[:] = tmp
+{% endif %}
+
 func_type = numba.types.FunctionType(
     {{ func_type }}
 )

sympleints/templates/py_equi_func.tpl

@@ -3,5 +3,4 @@ def {{ equi_name }}({{ args }}, result):
 
     # Calculate values w/ swapped arguments
     {{ name }}({{ equi_args }}, result)
-    # Swap two axes
-    result[:] = numpy.moveaxis(result.reshape({{ shape|join(",") }}), {{ from_axes }}, {{ to_axes }}).flatten()
+    {{ resort_func }}(result, {{ sizes|join(", ")}}, {{ ncomponents }})

sympleints/templates/py_module.tpl

@@ -20,6 +20,23 @@ from {{ boys }} import boys
 {{ ai }}
 {% endfor %}
 
+{% if resort_func == "resort_ba_ab" %}
+def resort_ba_ab(result, sizea, sizeb, ncomponents):
+    assert ncomponents > 0
+    tmp = numpy.zeros_like(result)
+    i = 0  # Original index
+    sizeab = sizea * sizeb
+    component_size = 0
+    for _ in range(ncomponents):
+        for b in range(sizeb):
+            for a in range(sizea):
+                j = component_size + (a * sizeb) + b  # New index
+                tmp[j] = result[i]
+                i += 1
+        component_size += sizeab
+    result[:] = tmp
+{% endif %}
+
 {% for func in funcs %}
 {{ func.text }}
 {% endfor %}