From 6515b7bf338cc9513e7f0903edfd6aba78b78961 Mon Sep 17 00:00:00 2001
From: Vonka Jakub <jakub.vonka@psi.ch>
Date: Wed, 13 Jul 2022 16:37:00 +0200
Subject: [PATCH] added small things to utilities

---
 src/cristallina/utils.py | 21 ++++++++++++++++++---
 tests/test_utils.py      | 12 +++++++++---
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/src/cristallina/utils.py b/src/cristallina/utils.py
index 9205cf7..0ba74e2 100644
--- a/src/cristallina/utils.py
+++ b/src/cristallina/utils.py
@@ -1,7 +1,8 @@
 import yaml
 import os
-
+import numpy as np
 from sfdata import SFDataFiles, sfdatafile, SFScanInfo
+from xraydb import material_mu
 
 
 def print_run_info(
@@ -136,7 +137,7 @@ def heuristic_extract_smalldata_path():
 
 ######################## Little useful functions ########################
 
- def find_nearest(array, value):
+def find_nearest(array, value):
     '''Finds an index in an array with a value that is nearest to given number'''
     array = np.asarray(array)
     idx = (np.abs(array - value)).argmin()
@@ -154,7 +155,7 @@ def gauss(x, H, A, x0, sigma):
     """Returns gauss function value"""
     return H + A * np.exp(-(x - x0) ** 2 / (2 * sigma ** 2))
 
-def gauss_fit(x, y, fit_details=None, plot=None:
+def gauss_fit(x, y, fit_details=None, plot=None):
     '''Returns [baseline_offset, Amplitude, center, sigma, FWHM]'''
 
     # Initial guesses
@@ -191,6 +192,20 @@ def gauss_fit(x, y, fit_details=None, plot=None:
 
     return popt
 
+def xray_transmission(energy,thickness,material='Si',density=[]):
+    '''Calculate x-ray tranmission for given energy, thickness and material. Default material is Si. Add material=element as a string for another material. Density as optional parameter'''
+    
+    mu = material_mu(material, energy)
+
+    if density == []:
+        mu_array = material_mu(material, energy)
+    else:
+        mu_array = material_mu(formula, energy, density=density)
+
+    trans = np.exp(-0.1*(thickness*1000)*mu_array) # 0.1 is beccause mu is in 1/cm, thickness converted to mm from m
+
+    return trans    
+
 ######################## Unit conversions ########################
 
 def joules_to_eV(joules):
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 48386f9..cabf9c5 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -1,7 +1,7 @@
 import pytest
 import os
-
-from cristallina.utils import ROI, print_run_info, heuristic_extract_pgroup, gauss
+import numpy as np
+from cristallina.utils import ROI, print_run_info, heuristic_extract_pgroup, gauss, find_nearest, xray_transmission
 
 __author__ = "Alexander Steppke"
 
@@ -33,4 +33,10 @@ def test_gauss():
 def test_find_nearest():
     # Make an array of 1 to 100 and check that the nearest index is the value
     a = np.linspace(0,99,100)
-    assert find_nearest(a,10.1) == 10
\ No newline at end of file
+    assert find_nearest(a,10.1) == 10
+
+def test_xray_transmission():
+    T = xray_transmission(9000,100e-6,material='Si')
+    assert T == 0.342385039732607
+
+