From a0114aaee449b8d7bf01a7720ce95b9e960e4f92 Mon Sep 17 00:00:00 2001
From: KuangYu <windwhisper.yu@gmail.com>
Date: Mon, 25 Apr 2022 17:19:29 +0800
Subject: [PATCH] Add convert_harm2cart in multipole.py

---
 dmff/admp/multipole.py | 43 +++++++++++++++++++++++++++++++++++++++---
 1 file changed, 40 insertions(+), 3 deletions(-)

diff --git a/dmff/admp/multipole.py b/dmff/admp/multipole.py
index c894db6fa..7f957a823 100644
--- a/dmff/admp/multipole.py
+++ b/dmff/admp/multipole.py
@@ -50,9 +50,6 @@ def convert_cart2harm(Theta, lmax):
     if lmax > 2:
         sys.exit('l > 2 (beyond quadrupole) not supported')
 
-    # n_sites = Theta.shape[0]
-    n_harm = (lmax + 1)**2
-    # Q = jnp.zeros((n_sites, n_harm))
     Q_mono = Theta[0:1]
     
     # dipole
@@ -74,6 +71,46 @@ def convert_cart2harm(Theta, lmax):
     return Q
 
 
+@partial(vmap, in_axes=(0, None), out_axes=0)
+@jit_condition(static_argnums=(1))
+def convert_harm2cart(Q, lmax):
+    '''
+    Convert the multipole moments in spherical representation to cartesian
+    Basically the inverse operation of convert_cart2harm
+
+    Inputs:
+        Q:
+            n * N_harm: stores the spherical harmonics moments of each site
+        lmax:
+            integer, the maximum L, currently only supports up to quadrupole
+
+    Outputs:
+        Theta:
+            n * n_cart, stores the cartesian multipoles
+    '''
+
+    if lmax > 2:
+        sys.exit('l > 2 (beyond quadrupole) not supported')
+
+    T_mono = Q[0:1]
+
+    # dipole
+    if lmax >= 1:
+        T_dip = C1_h2c.dot(Q[1:4].T).T
+    # quadrupole
+    if lmax >= 2:
+        T_quad = C2_h2c.dot(Q[4:9].T).T
+
+    if lmax == 0:
+        T = T_mono
+    elif lmax == 1:
+        T = jnp.hstack([T_mono, T_dip])
+    else:
+        T = jnp.hstack([T_mono, T_dip, T_quad])
+
+    return T
+
+
 @partial(vmap, in_axes=(0, 0), out_axes=0)
 @jit_condition(static_argnums=())
 def rot_ind_global2local(U_g, localframes):