From e02899881ae8acfbffe40ed2c2124bb62caac45b Mon Sep 17 00:00:00 2001
From: Srinath Kailasa <srinathkailasa@gmail.com>
Date: Wed, 15 Nov 2023 16:45:51 +0000
Subject: [PATCH] Start work on linear data structure

---
 fmm/src/fmm.rs   | 82 +++++++++++++++++++++++++++++++++++++++++++++++-
 fmm/src/types.rs | 27 ++++++++++++++++
 2 files changed, 108 insertions(+), 1 deletion(-)

diff --git a/fmm/src/fmm.rs b/fmm/src/fmm.rs
index 2c53873a..bf46dc32 100644
--- a/fmm/src/fmm.rs
+++ b/fmm/src/fmm.rs
@@ -27,7 +27,7 @@ use bempp_traits::{
 use bempp_tree::{constants::ROOT, types::single_node::SingleNodeTree};
 
 use crate::pinv::{pinv, SvdScalar};
-use crate::types::{C2EType, ChargeDict, FmmData, KiFmm};
+use crate::types::{C2EType, ChargeDict, FmmData, FmmDataLinear, KiFmm};
 
 /// Implementation of constructor for single node KiFMM
 impl<'a, T, U, V> KiFmm<SingleNodeTree<V>, T, U, V>
@@ -310,6 +310,86 @@ where
     }
 }
 
+/// Implementation of the data structure to store the data for the single node KiFMM.
+impl<T, U, V> FmmDataLinear<KiFmm<SingleNodeTree<V>, T, U, V>, V>
+where
+    T: Kernel<T = V>,
+    U: FieldTranslationData<T>,
+    V: Float + Scalar<Real = V> + Default,
+{
+    /// Constructor fo the KiFMM's associated FmmData on a single node.
+    ///
+    /// # Arguments
+    /// `fmm` - A single node KiFMM object.
+    /// `global_charges` - The charge data associated to the point data via unique global indices.
+    pub fn new(
+        fmm: KiFmm<SingleNodeTree<V>, T, U, V>,
+        global_charges: &ChargeDict<V>,
+    ) -> Result<Self, String> {
+        if let Some(keys) = fmm.tree().get_all_keys() {
+            let ncoeffs = fmm.m2l.ncoeffs(fmm.order);
+            let nkeys = keys.len();
+            let mut npoints = 0;
+
+            for key in keys.iter() {
+                if let Some(points) = fmm.tree().get_points(key) {
+                    npoints += points.len();
+                }
+            }
+
+            let multipoles = vec![V::default(); ncoeffs * nkeys];
+            let locals = vec![V::default(); ncoeffs * nkeys];
+
+            let potentials = vec![V::default(); npoints];
+            let points = vec![V::default(); npoints * 3];
+            let charges = vec![V::default(); npoints];
+            let global_indices = vec![0usize; npoints];
+
+            // Assign point coordinates and associated charges.
+            fmm.tree().get_all_points();
+
+            return Ok(Self {
+                fmm,
+                multipoles,
+                locals,
+                potentials,
+                points,
+                charges,
+                global_indices,
+            });
+        }
+
+        Err("Not a valid tree".to_string())
+
+        // let dummy = rlst_col_vec![V, ncoeffs];
+
+        // if let Some(keys) = fmm.tree().get_all_keys() {
+        //     for key in keys.iter() {
+        //         multipoles.insert(*key, Arc::new(Mutex::new(dummy.new_like_self().eval())));
+        //         locals.insert(*key, Arc::new(Mutex::new(dummy.new_like_self().eval())));
+        //         if let Some(point_data) = fmm.tree().get_points(key) {
+        //             points.insert(*key, point_data.iter().cloned().collect_vec());
+
+        //             let npoints = point_data.len();
+        //             potentials.insert(*key, Arc::new(Mutex::new(rlst_col_vec![V, npoints])));
+
+        //             // Lookup indices and store with charges
+        //             let mut tmp_idx = Vec::new();
+        //             for point in point_data.iter() {
+        //                 tmp_idx.push(point.global_idx)
+        //             }
+        //             let mut tmp_charges = vec![V::zero(); point_data.len()];
+        //             for i in 0..tmp_idx.len() {
+        //                 tmp_charges[i] = *global_charges.get(&tmp_idx[i]).unwrap();
+        //             }
+
+        //             charges.insert(*key, Arc::new(tmp_charges));
+        //         }
+        //     }
+        // }
+    }
+}
+
 impl<T, U, V, W> Fmm for KiFmm<T, U, V, W>
 where
     T: Tree,
diff --git a/fmm/src/types.rs b/fmm/src/types.rs
index 8641103a..9736d6c4 100644
--- a/fmm/src/types.rs
+++ b/fmm/src/types.rs
@@ -54,6 +54,33 @@ where
     pub charges: HashMap<MortonKey, Arc<Vec<U>>>,
 }
 
+pub struct FmmDataLinear<T, U>
+where
+    T: Fmm,
+    U: Scalar<Real = U> + Float + Default,
+{
+    /// The associated FMM object, which implements an FMM interface
+    pub fmm: T,
+
+    /// The multipole expansion data at each box.
+    pub multipoles: Vec<U>,
+
+    /// The local expansion data at each box.
+    pub locals: Vec<U>,
+
+    /// The evaluated potentials at each leaf box.
+    pub potentials: Vec<U>,
+
+    /// The point data at each leaf box. (matches charges/global idxs, but displaced by factor 3)
+    pub points: Vec<U>,
+
+    /// The charge data at each leaf box.
+    pub charges: Vec<U>,
+
+    /// Global indices of each charge
+    pub global_indices: Vec<usize>,
+}
+
 /// Type to store data associated with the kernel independent (KiFMM) in.
 pub struct KiFmm<T, U, V, W>
 where