added unit test for parray generateion with different current shape. added unit test for get_envelope test for slice with Imax

Author: sergey-tomin

ocelot/cpbd/beam.py

@@ -1117,20 +1117,37 @@ def recalculate_ref_particle(p_array):
     return p_array
 
 
-def get_envelope(p_array, tws_i=None, bounds=None):
+def get_envelope(p_array, tws_i=None, bounds=None, slice=None):
     """
-    Function to calculate twiss parameters form the ParticleArray
+    Calculate Twiss parameters from a ParticleArray.
+
+    This function processes particle data to compute Twiss parameters, optionally considering bounds
+    based on standard deviations of the particle distribution and a reference slice for calculations.
 
     :param p_array: ParticleArray
-    :param tws_i: optional, design Twiss for dispersion correction.
-    :param bounds: optional, [left_bound, right_bound] - bounds in units of std(p_array.tau())
-    :return: Twiss()
+        The input particle array containing particle properties such as position and momentum.
+    :param tws_i: Twiss (optional)
+        Design Twiss parameters for dispersion correction. Defaults to None.
+    :param bounds: list (optional)
+        A list specifying bounds as `[left_bound, right_bound]` in units of std(p_array.tau()).
+        If provided, only particles within these bounds are considered.
+    :param slice: str or None (optional)
+        Reference slice definition when `bounds` is specified:
+        - None (default): Takes the central slice with z0 = np.mean(tau).
+        - "Imax": Defines the reference slice at the position of maximum beam current.
+    :return: Twiss
+        Computed Twiss parameters for the filtered particle array.
     """
 
+    tau = p_array.tau()
     if bounds is not None:
-        tau = p_array.tau()
-        z0 = np.mean(tau)
         sig0 = np.std(tau)
+        if slice == "Imax":
+            charge = np.sum(p_array.q_array)
+            B = s_to_cur(tau, 0.01 * sig0, charge, speed_of_light)
+            z0 = B[np.argmax(B[:, 1]), 0]
+        else:
+            z0 = np.mean(tau)
         inds = np.argwhere((z0 + sig0 * bounds[0] <= tau) * (tau <= z0 + sig0 * bounds[1]))
         p = p_array.p()[inds]
         x = p_array.x()[inds]
@@ -1144,7 +1161,6 @@ def get_envelope(p_array, tws_i=None, bounds=None):
         px = p_array.px()
         y = p_array.y()
         py = p_array.py()
-        tau = p_array.tau()
 
     tws = Twiss()
     tws.E = np.copy(p_array.E)

unit_tests/ebeam_test/beam/beam_conf.py

@@ -0,0 +1,9 @@
+"""Test parameters description file"""
+
+import pytest
+import numpy as np
+
+from ocelot import *
+from ocelot.cpbd.r_matrix import rot_mtx
+from ocelot.utils.acc_utils import *
+

unit_tests/ebeam_test/beam/beam_test.py

@@ -0,0 +1,135 @@
+"""Test of the demo file demos/ebeam/csr_ex.py"""
+
+import os
+import sys
+import copy
+import time
+
+FILE_DIR = os.path.dirname(os.path.abspath(__file__))
+REF_RES_DIR = FILE_DIR + '/ref_results/'
+
+from unit_tests.params import *
+from beam_conf import *
+
+
+@pytest.mark.parametrize('parameter', [0, 1, 2])
+def test_generate_parray_shape(parameter, update_ref_values=False):
+    """R matrix calculation test"""
+    if parameter == 0:
+        shape = "gauss"
+    elif parameter == 1:
+        shape = "tri"
+    else:
+        shape = "rect"
+
+    parray = generate_parray(sigma_x=1e-4, sigma_px=2e-5, sigma_y=None, sigma_py=None,
+                    sigma_tau=1e-3, sigma_p=1e-4, chirp=0.0, charge=0.25e-9, nparticles=5000000, energy=0.13,
+                    tau_trunc=2.5, tws=None, shape=shape)
+
+    s, I = get_current(parray, num_bins=100)
+    B = np.column_stack((s, I))
+    if update_ref_values:
+        return numpy2json(B)
+    B_ref = json2numpy(json_read(REF_RES_DIR + sys._getframe().f_code.co_name + str(parameter) + '.json'))
+    result = check_matrix(B, B_ref, tolerance=1.0e-1, tolerance_type='relative', assert_info=' current - ')
+
+    assert check_result(result)
+
+
+def test_generate_parray_arb_shape(parameter=None, update_ref_values=False):
+    sigma_tau = 30e-6
+    A = np.array([0.2, 0.9, 0.3])
+
+    mu = (np.array([0.1, 0.3, 0.2]) - 0.5) * 2 * sigma_tau
+
+    sigma = np.array([0.9, 0.8, 0.7]) * sigma_tau
+
+    s = np.linspace(-2 * sigma_tau, 2 * sigma_tau, num=500)
+
+    s_reshaped = s[:, np.newaxis]
+
+    f = np.sum(A * np.exp(-(s_reshaped - mu) ** 2 / (2.0 * sigma ** 2)), axis=1)
+
+    shape = [s, f]
+
+    parray = generate_parray(sigma_tau=sigma_tau, tau_trunc=2.5, sigma_p=2.5/14000, chirp=-0.00, charge=250e-12, nparticles=5000000, energy=14, tws=None, shape=shape)
+
+    s, I = get_current(parray, num_bins=100)
+    B = np.column_stack((s, I))
+    if update_ref_values:
+        return numpy2json(B)
+    B_ref = json2numpy(json_read(REF_RES_DIR + sys._getframe().f_code.co_name + '.json'))
+    result = check_matrix(B, B_ref, tolerance=1.0e-1, tolerance_type='relative', assert_info=' current - ')
+
+    assert check_result(result)
+
+def test_get_envelope_slice(parameter=None, update_ref_values=False):
+    sigma_x = 1e-4
+    sigma_px = 2e-5
+    sigma_y = 5e-5
+    sigma_py = 4e-5
+    sigma_p = 1e-4
+    parray = generate_parray(sigma_x=sigma_x, sigma_px=sigma_px, sigma_y=sigma_y, sigma_py=sigma_py,
+                    sigma_tau=1e-3, sigma_p=sigma_p, chirp=0.0, charge=0.25e-9, nparticles=5000000, energy=0.13,
+                    tau_trunc=2.5, tws=None, shape="gauss")
+
+    tws = get_envelope(parray, bounds=[-0.5, 0.5], slice="Imax")
+
+    assert np.isclose(sigma_x, np.sqrt(tws.xx), rtol=1e-05, atol=1e-07, equal_nan=False)
+    assert np.isclose(sigma_px, np.sqrt(tws.pxpx), rtol=1e-05, atol=1e-07, equal_nan=False)
+    assert np.isclose(sigma_y, np.sqrt(tws.yy), rtol=1e-05, atol=1e-07, equal_nan=False)
+    assert np.isclose(sigma_py, np.sqrt(tws.pypy), rtol=1e-05, atol=1e-07, equal_nan=False)
+    assert np.isclose(sigma_p, np.sqrt(tws.pp), rtol=1e-05, atol=1e-07, equal_nan=False)
+
+
+def setup_module(module):
+
+    f = open(pytest.TEST_RESULTS_FILE, 'a')
+    f.write('### CSR_EX START ###\n\n')
+    f.close()
+
+
+def teardown_module(module):
+
+    f = open(pytest.TEST_RESULTS_FILE, 'a')
+    f.write('### CSR_EX END ###\n\n\n')
+    f.close()
+
+
+def setup_function(function):
+    
+    f = open(pytest.TEST_RESULTS_FILE, 'a')
+    f.write(function.__name__)
+    f.close()
+
+    pytest.t_start = time.time()
+
+
+def teardown_function(function):
+    f = open(pytest.TEST_RESULTS_FILE, 'a')
+    f.write(' execution time is ' + '{:.3f}'.format(time.time() - pytest.t_start) + ' sec\n\n')
+    f.close()
+    
+
+@pytest.mark.update
+def test_update_ref_values(cmdopt):
+
+    update_functions = []
+    update_functions.append('test_generate_parray_shape')
+    update_functions.append('test_generate_parray_arb_shape')
+
+    update_function_parameters = {}
+    update_function_parameters['test_generate_parray_shape'] = [0, 1, 2]
+
+    parameter = update_function_parameters[cmdopt] if cmdopt in update_function_parameters.keys() else ['']
+    print(parameter)
+    if cmdopt in update_functions:
+        for p in parameter:
+            print(p)
+            result = eval(cmdopt)( p, True)
+        
+            if os.path.isfile(REF_RES_DIR + cmdopt + str(p) + '.json'):
+                os.rename(REF_RES_DIR + cmdopt + str(p) + '.json', REF_RES_DIR + cmdopt + str(p) + '.old')
+            
+            json_save(result, REF_RES_DIR + cmdopt + str(p) + '.json')
+

unit_tests/ebeam_test/beam/ref_results/test_generate_parray_arb_shape.json

@@ -0,0 +1 @@
+[{"0": -5.939975583547192e-05, "1": 238.42609277358576}, {"0": -5.8199774201901335e-05, "1": 259.2120213309938}, {"0": -5.699979256833075e-05, "1": 281.3220535104485}, {"0": -5.579981093476015e-05, "1": 308.22883843392026}, {"0": -5.4599829301189566e-05, "1": 330.9509562499813}, {"0": -5.3399847667618974e-05, "1": 359.05692935945746}, {"0": -5.219986603404838e-05, "1": 387.42522488462197}, {"0": -5.0999884400477796e-05, "1": 419.66589892709214}, {"0": -4.97999027669072e-05, "1": 451.5568097486444}, {"0": -4.859992113333661e-05, "1": 483.69755144228094}, {"0": -4.739993949976603e-05, "1": 517.5871092405071}, {"0": -4.619995786619543e-05, "1": 546.0428455709006}, {"0": -4.499997623262484e-05, "1": 586.5404299357527}, {"0": -4.379999459905426e-05, "1": 619.0184433067006}, {"0": -4.260001296548366e-05, "1": 660.1530963953669}, {"0": -4.1400031331913074e-05, "1": 698.3897113678564}, {"0": -4.020004969834249e-05, "1": 738.4625832501647}, {"0": -3.900006806477189e-05, "1": 774.31331339425}, {"0": -3.7800086431201305e-05, "1": 817.8838174857373}, {"0": -3.660010479763071e-05, "1": 846.1896552928791}, {"0": -3.540012316406012e-05, "1": 892.2459765616046}, {"0": -3.4200141530489535e-05, "1": 923.7371579878196}, {"0": -3.3000159896918937e-05, "1": 966.5706610066593}, {"0": -3.180017826334835e-05, "1": 999.5233530345671}, {"0": -3.0600196629777766e-05, "1": 1034.4996751263575}, {"0": -2.940021499620717e-05, "1": 1066.2406874246578}, {"0": -2.8200233362636582e-05, "1": 1096.6451045573062}, {"0": -2.7000251729065993e-05, "1": 1127.3992849108731}, {"0": -2.58002700954954e-05, "1": 1153.6440180233203}, {"0": -2.460028846192481e-05, "1": 1185.4349964960372}, {"0": -2.340030682835422e-05, "1": 1205.059211498253}, {"0": -2.2200325194783632e-05, "1": 1224.471070259196}, {"0": -2.100034356121304e-05, "1": 1252.8518573279632}, {"0": -1.980036192764245e-05, "1": 1253.8386892726958}, {"0": -1.8600380294071863e-05, "1": 1261.9956672462456}, {"0": -1.740039866050127e-05, "1": 1279.3464213124942}, {"0": -1.620041702693068e-05, "1": 1280.4706602368733}, {"0": -1.500043539336009e-05, "1": 1279.0591158095976}, {"0": -1.3800453759789502e-05, "1": 1285.5672100273919}, {"0": -1.260047212621891e-05, "1": 1273.650277428974}, {"0": -1.1400490492648321e-05, "1": 1265.0810785164863}, {"0": -1.0200508859077732e-05, "1": 1260.209376510843}, {"0": -9.00052722550714e-06, "1": 1242.009197479508}, {"0": -7.800545591936548e-06, "1": 1224.8208334801136}, {"0": -6.60056395836596e-06, "1": 1210.8677792742105}, {"0": -5.400582324795371e-06, "1": 1185.9471497838103}, {"0": -4.200600691224779e-06, "1": 1160.7517063341172}, {"0": -3.0006190576541903e-06, "1": 1140.1031847563563}, {"0": -1.8006374240836017e-06, "1": 1102.8034355541822}, {"0": -6.006557905130097e-07, "1": 1075.209615732481}, {"0": 5.993258430575789e-07, "1": 1040.982786256943}, {"0": 1.799307476628171e-06, "1": 1000.7475243077799}, {"0": 2.999289110198763e-06, "1": 972.8289243523686}, {"0": 4.199270743769348e-06, "1": 930.0204044207381}, {"0": 5.39925237733994e-06, "1": 892.3459089104381}, {"0": 6.599234010910532e-06, "1": 858.9310297691743}, {"0": 7.799215644481117e-06, "1": 819.7325659391607}, {"0": 8.99919727805171e-06, "1": 780.7589498940226}, {"0": 1.0199178911622302e-05, "1": 742.3224702238659}, {"0": 1.1399160545192894e-05, "1": 706.2843669257176}, {"0": 1.2599142178763486e-05, "1": 664.5126451132365}, {"0": 1.379912381233407e-05, "1": 630.3357818121159}, {"0": 1.4999105445904663e-05, "1": 589.9880959705146}, {"0": 1.6199087079475255e-05, "1": 550.3524281143533}, {"0": 1.739906871304584e-05, "1": 517.8619231997995}, {"0": 1.8599050346616432e-05, "1": 486.05845318347696}, {"0": 1.9799031980187024e-05, "1": 451.2695042457475}, {"0": 2.099901361375761e-05, "1": 424.71248254319363}, {"0": 2.21989952473282e-05, "1": 384.32732207077737}, {"0": 2.3398976880898793e-05, "1": 363.17913874884727}, {"0": 2.4598958514469385e-05, "1": 337.23420268289965}, {"0": 2.5798940148039977e-05, "1": 311.58906366345315}, {"0": 2.6998921781610563e-05, "1": 285.09449967892044}, {"0": 2.8198903415181155e-05, "1": 256.11411851714934}, {"0": 2.9398885048751747e-05, "1": 237.8140071369793}, {"0": 3.059886668232233e-05, "1": 215.3292286493984}, {"0": 3.1798848315892924e-05, "1": 195.61757284195238}, {"0": 3.2998829949463516e-05, "1": 178.7290058890584}, {"0": 3.41988115830341e-05, "1": 163.42686497389911}, {"0": 3.539879321660469e-05, "1": 147.41270607329986}, {"0": 3.6598774850175285e-05, "1": 132.46032837905858}, {"0": 3.779875648374587e-05, "1": 116.7334749813563}, {"0": 3.899873811731646e-05, "1": 110.05049915310326}, {"0": 4.0198719750887054e-05, "1": 95.62276629023903}, {"0": 4.1398701384457646e-05, "1": 85.84188764814147}, {"0": 4.259868301802824e-05, "1": 74.52454914272595}, {"0": 4.3798664651598824e-05, "1": 68.77843908478864}, {"0": 4.4998646285169416e-05, "1": 59.884460038590035}, {"0": 4.619862791874001e-05, "1": 53.81356984694326}, {"0": 4.739860955231059e-05, "1": 46.51850838208371}, {"0": 4.8598591185881185e-05, "1": 41.24707698110646}, {"0": 4.979857281945178e-05, "1": 37.81190248994829}, {"0": 5.099855445302236e-05, "1": 30.841621158798247}, {"0": 5.2198536086592954e-05, "1": 26.644462507783178}, {"0": 5.3398517720163546e-05, "1": 23.908814458460856}, {"0": 5.459849935373413e-05, "1": 21.89767593818281}, {"0": 5.5798480987304724e-05, "1": 18.887213929567828}, {"0": 5.6998462620875316e-05, "1": 16.251498229079193}, {"0": 5.819844425444591e-05, "1": 13.753189508236886}, {"0": 5.939842588801649e-05, "1": 12.154271926897813}]

unit_tests/ebeam_test/beam/ref_results/test_generate_parray_shape0.json

@@ -0,0 +1 @@
+[{"0": -0.0024749934219549974, "1": 1.423721175923843}, {"0": -0.002424993660515818, "1": 1.5886078145336786}, {"0": -0.002374993899076639, "1": 1.7930672464098751}, {"0": -0.0023249941376374595, "1": 2.0490912270876933}, {"0": -0.0022749943761982802, "1": 2.2634438572804787}, {"0": -0.0022249946147591005, "1": 2.565336302880668}, {"0": -0.002174994853319921, "1": 2.8786209162393543}, {"0": -0.002124995091880742, "1": 3.1820123312814523}, {"0": -0.0020749953304415626, "1": 3.5210792190409554}, {"0": -0.0020249955690023833, "1": 3.8844294117593603}, {"0": -0.001974995807563204, "1": 4.323627458238292}, {"0": -0.0019249960461240247, "1": 4.836275007370698}, {"0": -0.0018749962846848454, "1": 5.214015306731415}, {"0": -0.0018249965232456661, "1": 5.728761413082491}, {"0": -0.0017749967618064869, "1": 6.251002366643106}, {"0": -0.0017249970003673076, "1": 6.828705189554303}, {"0": -0.0016749972389281283, "1": 7.482855454002785}, {"0": -0.001624997477488949, "1": 8.101630039622249}, {"0": -0.0015749977160497695, "1": 8.710811220813506}, {"0": -0.0015249979546105902, "1": 9.476185017851531}, {"0": -0.0014749981931714109, "1": 10.090462695145273}, {"0": -0.0014249984317322316, "1": 10.946074452586144}, {"0": -0.001374998670293052, "1": 11.773505584519151}, {"0": -0.0013249989088538728, "1": 12.655798998025965}, {"0": -0.0012749991474146935, "1": 13.456848267781373}, {"0": -0.0012249993859755142, "1": 14.284878987491084}, {"0": -0.001174999624536335, "1": 15.07723423448341}, {"0": -0.0011249998630971556, "1": 16.11691943939046}, {"0": -0.0010750001016579763, "1": 17.148210415422835}, {"0": -0.0010250003402187968, "1": 17.954056387392328}, {"0": -0.0009750005787796175, "1": 18.88851393747747}, {"0": -0.0009250008173404382, "1": 19.61161679625366}, {"0": -0.0008750010559012589, "1": 20.575454147400194}, {"0": -0.0008250012944620796, "1": 21.582461818473675}, {"0": -0.0007750015330229003, "1": 22.427880583709513}, {"0": -0.0007250017715837209, "1": 23.313171936100115}, {"0": -0.0006750020101445415, "1": 24.149297090796114}, {"0": -0.0006250022487053622, "1": 24.914071300057337}, {"0": -0.0005750024872661829, "1": 25.681243860425603}, {"0": -0.0005250027258270036, "1": 26.331197010436714}, {"0": -0.00047500296438782435, "1": 27.187408355654288}, {"0": -0.00042500320294864505, "1": 27.632902073789186}, {"0": -0.00037500344150946576, "1": 28.191418087843942}, {"0": -0.00032500368007028647, "1": 28.807494528467927}, {"0": -0.00027500391863110696, "1": 29.16874616396764}, {"0": -0.00022500415719192745, "1": 29.49192397564292}, {"0": -0.00017500439575274815, "1": 29.68589062142575}, {"0": -0.00012500463431356886, "1": 30.040846585269442}, {"0": -7.500487287438956e-05, "1": 30.033351738059906}, {"0": -2.500511143521027e-05, "1": 30.227318383842743}, {"0": 2.4994650003969023e-05, "1": 30.306164176487073}, {"0": 7.499441144314832e-05, "1": 30.24680498658754}, {"0": 0.0001249941728823276, "1": 30.00217317366824}, {"0": 0.0001749939343215069, "1": 29.757841154637294}, {"0": 0.0002249936957606862, "1": 29.523702127811347}, {"0": 0.0002749934571998655, "1": 28.988270243161974}, {"0": 0.0003249932186390448, "1": 28.550870960013356}, {"0": 0.0003749929800782243, "1": 28.239085316096613}, {"0": 0.0004249927415174038, "1": 27.677571363158037}, {"0": 0.0004749925029565831, "1": 26.892411169486888}, {"0": 0.0005249922643957624, "1": 26.414839505295145}, {"0": 0.0005749920258349417, "1": 25.63657457105674}, {"0": 0.000624991787274121, "1": 24.91616985727601}, {"0": 0.0006749915487133003, "1": 24.160089670777857}, {"0": 0.0007249913101524796, "1": 23.267003677289345}, {"0": 0.0007749910715916589, "1": 22.561888451816053}, {"0": 0.0008249908330308382, "1": 21.56567336072431}, {"0": 0.0008749905944700175, "1": 20.67648468778477}, {"0": 0.0009249903559091967, "1": 19.659284024506324}, {"0": 0.000974990117348376, "1": 18.826756396470884}, {"0": 0.0010249898787875556, "1": 17.937867517419722}, {"0": 0.001074989640226735, "1": 17.055574103912964}, {"0": 0.0011249894016659144, "1": 16.081243966673064}, {"0": 0.0011749891631050936, "1": 15.19925034705468}, {"0": 0.001224988924544273, "1": 14.345737146832537}, {"0": 0.0012749886859834522, "1": 13.472437549977213}, {"0": 0.0013249884474226315, "1": 12.662094669681974}, {"0": 0.0013749882088618108, "1": 11.71474598239637}, {"0": 0.0014249879703009901, "1": 10.998538383052912}, {"0": 0.0014749877317401694, "1": 10.185197563873837}, {"0": 0.0015249874931793487, "1": 9.488776361163561}, {"0": 0.001574987254618528, "1": 8.830728776166113}, {"0": 0.0016249870160577073, "1": 7.948435362659294}, {"0": 0.0016749867774968866, "1": 7.364736661980464}, {"0": 0.0017249865389360659, "1": 6.910249127194075}, {"0": 0.0017749863003752452, "1": 6.286677839360506}, {"0": 0.0018249860618144245, "1": 5.693385734253472}, {"0": 0.0018749858232536038, "1": 5.2350008789181235}, {"0": 0.001924985584692783, "1": 4.733745497544215}, {"0": 0.001974985346131963, "1": 4.296646008283956}, {"0": 0.0020249851075711425, "1": 3.914708594485893}, {"0": 0.002074984869010322, "1": 3.519280455710663}, {"0": 0.002124984630449501, "1": 3.2179875978872357}, {"0": 0.0021749843918886804, "1": 2.7754918186361124}, {"0": 0.0022249841533278597, "1": 2.550346608461593}, {"0": 0.002274983914767039, "1": 2.274536231150595}, {"0": 0.0023249836762062183, "1": 2.0406969982130105}, {"0": 0.0023749834376453976, "1": 1.8419336502160635}, {"0": 0.002424983199084577, "1": 1.5556304868117117}, {"0": 0.0024749829605237557, "1": 1.418624679821357}]