Improve ParticleGroup type

.gitignore

@@ -5,3 +5,4 @@ deps/deps.jl
 Manifest.toml
 *.swp
 .ipynb_checkpoints/maxwell_1d_fem-checkpoint.ipynb
+.ipynb_checkpoints/particle_group-checkpoint.ipynb

particle_group.ipynb

@@ -0,0 +1,159 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mutable struct ParticleGroup{D, V}\n",
+    "    \n",
+    "    n    :: Int\n",
+    "    dims :: Tuple{Int, Int}\n",
+    "    data :: Array{Float64, 2}\n",
+    "    \n",
+    "    function ParticleGroup{D, V}( n ) where {D, V}\n",
+    "        \n",
+    "        dims = (D, V)\n",
+    "        data = zeros(Float64, (D+V,n))\n",
+    "        \n",
+    "        new(n, dims, data)\n",
+    "        \n",
+    "    end\n",
+    "    \n",
+    "end"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "get_x (generic function with 1 method)"
+      ]
+     },
+     "execution_count": 19,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "@generated function get_x( p ::ParticleGroup{D,V}, i) where {D, V}\n",
+    "    \n",
+    "    return :(p.data[1:$D, i])\n",
+    "           \n",
+    "end"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "get_v (generic function with 1 method)"
+      ]
+     },
+     "execution_count": 20,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "@generated function get_v( p ::ParticleGroup{D,V}, i) where {D, V}\n",
+    "    \n",
+    "    return :(p.data[$D+1:$D+$V, i])\n",
+    "           \n",
+    "end"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "ParticleGroup{1,2}(5, (1, 2), [0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0])"
+      ]
+     },
+     "execution_count": 21,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "p1d1v = ParticleGroup{1,1}(5)\n",
+    "p1d2v = ParticleGroup{1,2}(5)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "([0.0], [0.0, 0.0])"
+      ]
+     },
+     "execution_count": 26,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "get_x(p1d2v, 1), get_v(p1d2v, 1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "([0.0], [0.0])"
+      ]
+     },
+     "execution_count": 27,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "get_x(p1d1v, 1), get_v(p1d1v, 1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Julia 1.1.0",
+   "language": "julia",
+   "name": "julia-1.1"
+  },
+  "language_info": {
+   "file_extension": ".jl",
+   "mimetype": "application/julia",
+   "name": "julia",
+   "version": "1.1.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

src/particle_group.jl

@@ -1,7 +1,10 @@
-export ParticleGroup1D2V
+export ParticleGroup
 
-mutable struct ParticleGroup1D2V
+abstract type AbstractParticleGroup end
 
+mutable struct ParticleGroup{D,V} <:  AbstractParticleGroup 
+
+    dims              :: Tuple{Int64, Int64}
     n_particles       :: Int64
     n_total_particles :: Int64
     particle_array    :: Array{Float64, 2} 
@@ -10,16 +13,18 @@ mutable struct ParticleGroup1D2V
     mass              :: Float64
     n_weights         :: Int64
 
-    function ParticleGroup1D2V( n_particles, 
-                            n_total_particles, 
-                            charge, 
-                            mass, 
-                            n_weights)
+    function ParticleGroup{D,V}( n_particles, 
+                                 n_total_particles, 
+                                 charge, 
+                                 mass, 
+                                 n_weights) where {D, V}
 
-        particle_array = zeros( Float64, (3+n_weights, n_particles)) 
+        dims = (D, V)
+        particle_array = zeros( Float64, (sum(dims)+n_weights, n_particles)) 
         common_weight  = 1.0
 
-        new( n_particles,
+        new( dims,
+             n_particles,
              n_total_particles,
              particle_array,
              common_weight,
@@ -33,82 +38,98 @@ end
 """  
 Get position
 """
-function get_x( self :: ParticleGroup1D2V, i )
+@generated function get_x( p :: ParticleGroup{D,V}, i ) where {D, V}
 
-    self.particle_array[1, i]
+    if D == 1
+        return :(p.particle_array[1, i])
+    else
+        return :(p.particle_array[1:$D, i])
+    end
 
 end 
 
 """  
 Get velocities
 """
-function get_v( self :: ParticleGroup1D2V, i )
+@generated function get_v( p :: ParticleGroup{D,V}, i ) where {D, V}
 
-    self.particle_array[2:3, i]
+    if V == 1
+        return :(p.particle_array[$D+1, i])
+    else
+        return :(p.particle_array[$D+1:$D+$V, i])
+    end
 
 end
 
 
 """
 Get charge of particle (q * particle_weight)
 """
-function get_charge( self :: ParticleGroup1D2V, i; i_wi=1)
+@generated function get_charge( p :: ParticleGroup{D,V}, i; i_wi=1) where {D, V}
 
-    self.charge * self.particle_array[3+i_wi, i] * self.common_weight
+    return :(p.charge * p.particle_array[$D+$V+i_wi, i] * p.common_weight)
 
 end 
 
 
 """
 Get mass of particle (m * particle_weight)
 """
-function get_mass( self :: ParticleGroup1D2V, i; i_wi=1) 
+@generated function get_mass( p :: ParticleGroup{D,V}, i; i_wi=1) where {D,V}
 
-    self.mass * self.particle_array[3+i_wi, i] * self.common_weight
+    return :(p.mass * p.particle_array[$D+$V+i_wi, i] * p.common_weight)
 
 end
 
 """
 Get particle weights
 """
-function get_weights( self :: ParticleGroup1D2V, i) 
+@generated function get_weights( p :: ParticleGroup{D,V}, i) where {D, V}
 
-    self.particle_array[4:3+self.n_weights, i]
+    return :(p.particle_array[$D+$V+1:$D+$V+p.n_weights, i])
 
 end
 
 """
 Set position of particle @ i
 """
-function set_x( self :: ParticleGroup1D2V, i, x )
+@generated function set_x( p :: ParticleGroup{D,V}, i, x ) where {D, V}
 
-    self.particle_array[1, i] = x
+    if D == 1
+        return :(p.particle_array[1, i] = x)
+    else
+        return :(p.particle_array[1:$D, i] .= x)
+    end
 
 end
 
 """
 Set velocity of particle @ i
 """
-function set_v( self :: ParticleGroup1D2V, i, x )
+@generated function set_v( p :: ParticleGroup{D,V}, i, x ) where {D, V}
 
-    self.particle_array[2:3, i] .= x
+    if V == 1
+        return :(p.particle_array[$D+1, i] = x)
+    else
+        return :(p.particle_array[$D+1:$D+$V, i] .= x)
+    end
 
 end
 
 """
 Set weights of particle @ i
 """
-function set_weights( self :: ParticleGroup1D2V, i, x )
+@generated function set_weights( p :: ParticleGroup{D,V}, i, x ) where {D, V}
 
-    self.particle_array[4:3+self.n_weights, i] .= x
+    return :(p.particle_array[$D+$V+1:$D+$V+p.n_weights, i] .= x)
 
 end
 
 """
 Set the common weight
 """
-function set_common_weight( self :: ParticleGroup1D2V, x )
+function set_common_weight( p :: AbstractParticleGroup, x ) 
 
-    self.common_weight = x
+    p.common_weight = x
 
 end

test/test_particle_mesh_coupling_spline_1d.jl

@@ -11,7 +11,7 @@ x_vec = [0.1, 0.65, 0.7, 1.5] # Particle positions
 v_vec = [1.5  -3.0  0.0  6.0; 
          0.0   0.5  0.0  0.0]'
 
-particle_group = ParticleGroup1D2V( n_particles, n_particles ,1.0, 1.0, 1)
+particle_group = ParticleGroup{1,2}( n_particles, n_particles ,1.0, 1.0, 1)
 
 set_common_weight(particle_group, 1.0/n_particles)