From 8f63d653145d133b6641931e3a704682ab1098bc Mon Sep 17 00:00:00 2001
From: hrocho <hrocho@vodacionline.sk>
Date: Tue, 1 Aug 2017 17:34:44 +0200
Subject: [PATCH] python update 2

---
 catima.pyx    |  85 ++++++++++++++++++++++++++++++++++------
 tests/test.py | 106 +++++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 177 insertions(+), 14 deletions(-)

diff --git a/catima.pyx b/catima.pyx
index b9e40dd..b9311d8 100644
--- a/catima.pyx
+++ b/catima.pyx
@@ -12,16 +12,30 @@ import numpy
 
 cdef class Material:
     cdef catimac.Material cbase
-    def __cinit__(self, elements=None):
+    def __cinit__(self, elements=None, thickness=None, density=None):
         self.cbase = catimac.Material()
         if(elements and (isinstance(elements[0],float) or isinstance(elements[0],int))):
             self.cbase.add_element(elements[0],elements[1],elements[2])
         if(elements and isinstance(elements[0],list)):
             for e in elements:
                 self.cbase.add_element(e[0],e[1],e[2])
+        if(not thickness is None):
+            self.thickness(thickness)
+        if(not density is None):
+            self.density(density)
 
     cdef from_c(self, catimac.Material &other):
         self.cbase = other
+    
+    cdef catimac.Material getc(self):
+        cdef catimac.Material res
+        res = self.cbase 
+        return res
+    
+    def copy(self):
+        res = Material()
+        res.cbase = self.cbase
+        return res
 
     def add_element(self, a, z , s):
         self.cbase.add_element(a, z, s)
@@ -89,29 +103,42 @@ cdef class Target:
         return self.cbase.Z
 
 cdef class Layers:
-    cdef catimac.Layers cbase
     cdef public:
         materials
     def __init__(self):
         self.materials=[]
         
     def add(self,Material m):
-        self.cbase.add(m.cbase)
-        self.materials.append(m)
+        self.materials.append(m.copy())
         
     def num(self):
-        return self.cbase.num()
+        return len(self.materials)
 
     def get(self, key):
-        cdef catimac.Material cmat = self.cbase[key]
-        res = Material()
-        res.from_c(cmat)
-        return res
+        return self.materials[key]
         
     def __getitem__(self, key):
         if(isinstance(key,int) and key<self.num()):
             return self.get(key)
         return None
+    
+    def __add__(self, other):
+        res = Layers()
+        for e in self.materials:
+            res.add(e)
+            
+        if(isinstance(other,Layers)):
+            for e in other.materials:
+                res.add(e)
+        if(isinstance(other,Material)):
+            res.add(other.copy())
+        return res
+    
+    cdef catimac.Layers getc(self):
+        cdef catimac.Layers res
+        #for l in self.materials:
+        #    res.add(l.getc())
+        return res
 
 cdef class Projectile:
     cdef catimac.Projectile cbase
@@ -176,6 +203,11 @@ cdef class Result:
                 "tof":self.tof,
                 }
 
+    def __getitem__(self,key):
+        d = self.get_dict()
+        if(key in d):
+            return d[key]
+
     cdef setc(self,catimac.Result &val):
         self.Ein=val.Ein
         self.Eout=val.Eout
@@ -189,13 +221,27 @@ cdef class Result:
         self.tof=val.tof
 
 cdef class MultiResult:
+    cdef public Result total_result
+    cdef public results
+    cdef public total
+
     def __init__(self):
         self.total_result = Result()
         self.results = []
+        self.total = {}
+
     cdef setc(self, catimac.MultiResult &val):
         self.total_result.setc(val.total_result)
         for e in val.results:
             self.results.append(e)
+        self.total = self.total_result.get_dict()
+
+    def __getitem__(self,key):
+        if(isinstance(key,int) and key<len(self.results)):
+            return self.results[key]
+        if(isinstance(key,str) and key in self.total):
+            return self.total[key]
+        return None
 
 class z_eff_type(IntEnum):
     none = 0,
@@ -249,14 +295,29 @@ def calculate(Projectile projectile, Material material, energy = None, Config co
     res.setc(cres)
     return res
 
-def calculate(Projectile projectile, Layers layers, energy = None,  Config config = default_config):
+cdef catimac.Layers get_clayers(Layers layers):
+    cdef catimac.Layers res
+    cdef catimac.Material m
+    for l in layers.materials:
+        m = get_cmaterial(l)
+        res.add(m)
+    return res
+
+cdef catimac.Material get_cmaterial(Material material):
+    cdef catimac.Material res
+    res = material.cbase
+    return res
+
+def calculate_layers(Projectile projectile, Layers layers, energy = None,  Config config = default_config):
+    cdef catimac.Layers clayers
+    clayers = catimac.Layers()
+    clayers = get_clayers(layers)
     if(not energy is None):
         projectile.T(energy)
-    cdef catimac.MultiResult cres = catimac.calculate(projectile.cbase, layers.cbase, config.cbase)
+    cdef catimac.MultiResult cres = catimac.calculate(projectile.cbase, clayers, config.cbase)
     res = MultiResult()
     res.setc(cres)
     return res
-    
 
 def range(Projectile projectile, Material material, energy = None, Config config = default_config):  
     if(isinstance(energy,numpy.ndarray)):
diff --git a/tests/test.py b/tests/test.py
index 7f745b5..6b8ccd2 100644
--- a/tests/test.py
+++ b/tests/test.py
@@ -46,6 +46,22 @@ class TestStructures(unittest.TestCase):
         mat2 = catima.Material([0,6,1])
         self.assertEqual(mat2.ncomponents(),1)
         self.assertAlmostEqual(mat2.molar_mass(),12,1)
+
+        mat5 = catima.Material([0,6,1],density=1.9, thickness=0.5)
+        self.assertEqual(mat5.ncomponents(),1)
+        self.assertEqual(mat5.thickness(),0.5)
+        self.assertEqual(mat5.density(),1.9)
+
+        # copy
+        mat3.density(1.8)
+        matc = mat3.copy()
+        self.assertEqual(matc.ncomponents(),1)
+        self.assertEqual(matc.molar_mass(),12)
+        self.assertEqual(matc.density(),1.8)
+        mat3.density(2.0)
+        self.assertEqual(matc.density(),1.8)
+        self.assertEqual(mat3.density(),2.0)
+
     
     def test_default_material(self):
         m1 = catima.get_material(6);
@@ -59,10 +75,12 @@ class TestStructures(unittest.TestCase):
         self.assertAlmostEqual(m2.density(),1.0,1)
 
     def test_layers(self):
-        graphite = catima.get_material(6);
+        graphite = catima.get_material(6)
         graphite.thickness(0.5)
-        p10 = catima.get_material(catima.material.P10);
+        p10 = catima.get_material(catima.material.P10)
         p10.thickness(0.01)
+        n2 = catima.get_material(7)
+        n2.thickness(0.02)
 
         mat= catima.Layers()
         self.assertEqual(mat.num(),0)
@@ -84,8 +102,92 @@ class TestStructures(unittest.TestCase):
         self.assertAlmostEqual(mat[0].density(),2.0,1)
         self.assertAlmostEqual(mat[2].thickness(),1.0,1)
         self.assertAlmostEqual(mat[2].density(),1.8,1)
+        mat[2].thickness(1.2)
+        mat[2].density(1.9)
+        self.assertAlmostEqual(mat.materials[2].thickness(),1.2,1)
+        self.assertAlmostEqual(mat.materials[2].density(),1.9,1)
+        #self.assertAlmostEqual(mat.materials[0].thickness(),0.5,1)
+        #self.assertAlmostEqual(mat.materials[0].density(),2.0,1)
         self.assertEqual(mat[3],None)
         self.assertEqual(mat["a"],None)
 
+        mat2 = catima.Layers()
+        mat2.add(n2)
+        self.assertEqual(mat2.num(),1)
+
+        mats = mat2 + mat
+        self.assertEqual(mats.num(),4)
+        self.assertAlmostEqual(mats[0].molar_mass(),14,1)
+        self.assertEqual(mats[0].thickness(),0.02)
+        self.assertAlmostEqual(mats[1].molar_mass(),12,1)
+        self.assertAlmostEqual(mats[3].molar_mass(),12,1)
+
+        n2.thickness(0.5)
+        mats = mats + n2
+        self.assertEqual(mats.num(),5)
+        self.assertAlmostEqual(mats[0].molar_mass(),14,1)
+        self.assertEqual(mats[0].thickness(),0.02)
+        self.assertAlmostEqual(mats[4].molar_mass(),14,1)
+        self.assertEqual(mats[4].thickness(),0.5)
+
+    def test_material_calculation(self):
+        water = catima.get_material(catima.material.WATER)
+        p = catima.Projectile(1,1)
+        
+        p(1000)
+        res = catima.calculate(p,water)
+        res2 = catima.dedx(p,water)
+        self.assertAlmostEqual(res.dEdxi,2.23,1)
+        self.assertAlmostEqual(res["dEdxi"],2.23,1)
+        self.assertAlmostEqual(res.dEdxi,res2,3)
+        res = catima.calculate(p(500),water)
+        res2 = catima.dedx(p,water)
+        self.assertAlmostEqual(res.dEdxi,2.76,1)
+        self.assertAlmostEqual(res.dEdxi,res2,3)
+
+        res = catima.calculate(p(9),water)
+        res2 = catima.dedx(p,water)
+        self.assertAlmostEqual(res.dEdxi,51.17,1)
+        self.assertAlmostEqual(res.dEdxi,res2,3)
+        res = catima.calculate(p(9),water)
+        res = catima.calculate(p(9),water)
+        self.assertAlmostEqual(res.dEdxi,51.17,1)
+    
+    def test_eout(self):
+        graphite = catima.get_material(6)
+        graphite.thickness(0.5)
+        p = catima.Projectile(12,6)
+        res = catima.calculate(p(1000),graphite)
+        res2 = catima.energy_out(p(1000),graphite)
+        self.assertAlmostEqual(res.Eout,997.077,1)
+        self.assertAlmostEqual(res["Eout"],997.077,1)
+        self.assertAlmostEqual(res.Eout,res2,3)
+
+    def test_layer_calculateion(self):
+        p = catima.Projectile(12,6)
+        water = catima.get_material(catima.material.WATER)
+        water.thickness(10.0)
+        graphite = catima.get_material(6)
+        graphite.thickness(1.0)
+
+        mat = catima.Layers()
+        mat.add(water)
+        mat.add(graphite)
+        res = catima.calculate_layers(p(1000),mat)
+        self.assertEqual(len(res.results),2)
+        self.assertAlmostEqual(res.total_result.Eout,926.3,1)
+        self.assertAlmostEqual(res.total_result.sigma_a,0.00269,1)
+        self.assertAlmostEqual(res["Eout"],926.3,1)
+        self.assertAlmostEqual(res["sigma_a"],0.00269,4)
+        self.assertAlmostEqual(res["tof"],0.402,2)
+        self.assertAlmostEqual(res["Eloss"],884,0)
+
+        self.assertAlmostEqual(res[0]["Eout"],932,0)
+        self.assertAlmostEqual(res[1]["Eout"],926,0)
+        self.assertAlmostEqual(res[0]["sigma_a"],0.00258,4)
+        self.assertAlmostEqual(res[1]["sigma_a"],0.000774,4)
+        self.assertAlmostEqual(res[0]["range"],107.1,0)
+        self.assertAlmostEqual(res[1]["range"],110.7,0)
+
 if __name__ == "__main__":
     unittest.main()
\ No newline at end of file