Corrected executable and IO; works now.

freddieknets · freddieknets · commit ceb52255b7df · 2023-09-16T02:59:54.000+02:00
diff --git a/install_protection_hook.sh b/install_protection_hook.sh
@@ -18,7 +18,7 @@ then
   exit 1
 fi
 
-for file in pyproject.toml version.sh LICENSE install_protection_hook.sh pin_xsuite_versions.sh xboinc/general.py 'xboinc/executable/.*'
+for file in pyproject.toml version.sh LICENSE install_protection_hook.sh pin_xsuite_versions.sh xboinc/general.py 'xboinc/executable/.*' 'xboinc/simulation_io/.*'
 do
   git diff --name-only | grep '^'${file}'$' &> /dev/null
   if [ $? -eq 0 ]
diff --git a/xboinc/executable/__init__.py b/xboinc/executable/__init__.py
@@ -4,4 +4,3 @@
 # ######################################### #
 
 from .generate_executable_source import generate_executable_source, generate_executable
-from .default_tracker import get_default_tracker
diff --git a/xboinc/executable/generate_executable_source.py b/xboinc/executable/generate_executable_source.py
@@ -1,12 +1,23 @@
+# copyright ############################### #
+# This file is part of the Xboinc Package.  #
+# Copyright (c) CERN, 2023.                 #
+# ######################################### #
+
 import xobjects as xo
-from .default_tracker import get_default_tracker
-from .simulation_io import SimState, SimConfig
-from .general import _pkg_root, __version__
+from ..simulation_io import SimState, SimConfig, SimVersion, get_default_tracker
+from ..general import _pkg_root, __version__
 
 from pathlib import Path
 import shutil
 import subprocess
 
+# ===============================================================================================
+# IMPORTANT
+# ===============================================================================================
+# Only make changes to this file just before a minor version bump (need a separate commit though)
+# to avoid having multiple xboinc versions with out-of-sync executables.
+# ===============================================================================================
+
 insert_in_all_files = """
 #include <stdio.h>
 #ifndef NULL
@@ -21,6 +32,8 @@ def generate_executable_source(write_source_files=True,
 
     sim_config_sources = [
         insert_in_all_files,
+        xo.specialize_source(SimVersion._XoStruct._gen_c_api().source,
+                                    specialize_for='cpu_serial'),
         xo.specialize_source(SimState._XoStruct._gen_c_api().source,
                                     specialize_for='cpu_serial'),
         xo.specialize_source(SimConfig._gen_c_api().source,
@@ -34,7 +47,7 @@ def generate_executable_source(write_source_files=True,
     xtrack_tracker_h = (
         insert_in_all_files + track_kernel.specialized_source)
 
-    with open(_pkg_root.joinpath('executable_src/main.c'), 'r') as fid:
+    with open(_pkg_root.joinpath('executable/main.c'), 'r') as fid:
         main_c = insert_in_all_files + fid.read()
 
     dct_sources = {
@@ -44,7 +57,8 @@ def generate_executable_source(write_source_files=True,
     }
 
     for nn, vv in dct_sources.items():
-        with open(nn, 'w') as fid:
+        file = Path.cwd() / nn
+        with file.open('w') as fid:
             fid.write(vv)
 
     return dct_sources
@@ -64,17 +78,22 @@ def generate_executable(keep_source=False):
     else:
         raise RuntimeError("Neither clang or gcc are found. Install a C compiler.")
 
-    tag = '.'.join(__version__.split('.')[:2])
+    tag = '_'
+    tag += '.'.join(__version__.split('.')[:2])
     cmd = subprocess.run(['uname', '-ps'], stdout=subprocess.PIPE)
     if cmd.returncode == 0:
         tag += '-' + cmd.stdout.decode('UTF-8').strip().lower().replace(' ','-')
-    cmd = subprocess.run([compiler, 'main.c', '-o', f'xboinc_{tag}', '-lm'],
+    cmd = subprocess.run([compiler, 'main.c', '-o', f'xboinc{tag}', '-lm'],
                           stdout=subprocess.PIPE, stderr=subprocess.PIPE)
     if cmd.returncode != 0:
         stderr = cmd.stderr.decode('UTF-8').strip().split('\n')
         raise RuntimeError(f"Compilation failed. Stderr:\n {stderr}")
 
-    if not keep_source:
+    if keep_source:
+        main.rename(main.parent / f'{main.stem}{tag}{main.suffix}')
+        config.rename(config.parent / f'{config.stem}{tag}{config.suffix}')
+        tracker.rename(tracker.parent / f'{tracker.stem}{tag}{tracker.suffix}')
+    else:
         main.unlink()
         config.unlink()
         tracker.unlink()
diff --git a/xboinc/executable/main.c b/xboinc/executable/main.c
@@ -63,8 +63,7 @@ int main(){
 
     // Get sim config
     SimConfig sim_config = (SimConfig) sim_buffer;
-    SimVersion sim_ver = SimConfig_getp_sim_state_version(sim_config);
-    int64_t input_version = SimVersionData_get_xboinc_version(sim_ver);
+    int64_t input_version = SimConfig_get_sim_state_version_xboinc_version(sim_config);
     // Compatible if major and minor versions match (no new executables are made at patches)
     input_version = input_version/1000;
     if (input_version != xboinc_exec_version){
@@ -84,8 +83,9 @@ int main(){
     // Check if checkpoint exists
     printf("sim_state: %p\n", (int8_t*) sim_state);
     int8_t* loaded = file_to_buffer("./checkpoint.bin", (int8_t*) sim_state);
+    int64_t current_turn = SimStateData_get_i_turn(sim_state);
     if (loaded){
-       printf("Loaded checkpoint\n");
+       printf("Loaded checkpoint, continuing from turn %d\n", (int) current_turn);
     }
     else{
        printf("No checkpoint found\n");
@@ -98,7 +98,6 @@ int main(){
         step_turns = num_elements;
     }
 
-    int64_t current_turn = SimStateData_get_i_turn(sim_state);
     while (current_turn < num_turns){
         if (current_turn > num_turns - checkpoint_every){
             step_turns = num_turns - checkpoint_every;
@@ -121,16 +120,21 @@ int main(){
         );
         SimStateData_set_i_turn(sim_state, SimStateData_get_i_turn(sim_state) + step_turns);
 
+        current_turn = SimStateData_get_i_turn(sim_state);
         if (checkpoint_every>0){
-            printf("Checkpointing turn %d!\n", (int) SimStateData_get_i_turn(sim_state));
-            FILE *chkp_fid;
-            chkp_fid = fopen("./checkpoint.bin", "wb");
-            fwrite(SimConfig_getp_sim_state(sim_config), sizeof(int8_t),
-                SimConfig_get_sim_state_size(sim_config), chkp_fid);
-            fclose(chkp_fid);
+            if (current_turn < num_turns){
+                printf("Checkpointing turn %d\n", (int) current_turn);
+                FILE *chkp_fid;
+                chkp_fid = fopen("./checkpoint.bin", "wb");
+                fwrite(SimConfig_getp_sim_state(sim_config), sizeof(int8_t),
+                    SimConfig_get_sim_state_size(sim_config), chkp_fid);
+                fclose(chkp_fid);
+            }
         }
     }
 
+    printf("Finished tracking\n");
+
     // Quick check
     //for (int ii=0; ii<ParticlesData_get__capacity(particles); ii++){
         //printf("s[%d] = %e\n", ii, ParticlesData_get_s(particles, (int64_t) ii));
@@ -148,9 +152,6 @@ int main(){
         printf("Error: could not remove checkpoint file\n");
         return -1;  // error
     }
-    else{
-        printf("Checkpoint file removed\n");
-    }
 
     return 0;
 }
diff --git a/xboinc/simulation_io/__init__.py b/xboinc/simulation_io/__init__.py
@@ -3,5 +3,6 @@
 # Copyright (c) CERN, 2023.                 #
 # ######################################### #
 
-from .output import SimState
+from .default_tracker import get_default_tracker
+from .output import SimState, SimVersion
 from .input import SimConfig
diff --git a/xboinc/simulation_io/default_tracker.py b/xboinc/simulation_io/default_tracker.py
@@ -23,51 +23,80 @@ def get_default_tracker():
 
     # Dummy line containing all supported element types
     default_line = xt.Line(elements=[
-        xt.Drift(length=1.0),
-        xt.Multipole(knl=[0]),
+        xt.Marker(),
+        xt.Drift(),
+        xt.Bend(length=1.0),
+        xt.Multipole(),
+        xt.CombinedFunctionMagnet(length=1.0),
+        xt.Quadrupole(length=1.0),
+        xt.Sextupole(length=1.0),
         xt.SimpleThinBend(),
         xt.SimpleThinQuadrupole(),
         xt.ReferenceEnergyIncrease(),
         xt.Cavity(),
+        xt.Solenoid(length=1.0),
         xt.XYShift(),
         xt.Elens(),
+        xt.NonLinearLens(),
         xt.Wire(),
         xt.SRotation(),
         xt.XRotation(),
         xt.YRotation(),
         xt.ZetaShift(),
         xt.RFMultipole(knl=[0], pn=[0]),
+#        xt.Fringe(),
+#        xt.Wedge(),
         xt.DipoleEdge(),
         xt.Exciter(nsamples=1),
         xt.LineSegmentMap(),
-        # xt.EnergyChange(), # not working!!!
+        xt.FirstOrderTaylorMap(),
         xf.BeamBeamBiGaussian2D(
             other_beam_Sigma_11=1.,
             other_beam_Sigma_33=1.,
             other_beam_num_particles=0.,
             other_beam_q0=1.,
             other_beam_beta0=1.,
         ),
-        xf.BeamBeamBiGaussian3D(
-            slices_other_beam_zeta_center=[0],
-            slices_other_beam_num_particles=[0],
-            phi=0.,
-            alpha=0,
-            other_beam_q0=1.,
-            slices_other_beam_Sigma_11=[1],
-            slices_other_beam_Sigma_12=[0],
-            slices_other_beam_Sigma_22=[0],
-            slices_other_beam_Sigma_33=[1],
-            slices_other_beam_Sigma_34=[0],
-            slices_other_beam_Sigma_44=[0],
-        ),
+#         # Doesn't work because issue with definition of atomicAdd when generating executable
+#         xf.BeamBeamBiGaussian3D(
+#             slices_other_beam_zeta_center=[0],
+#             slices_other_beam_num_particles=[0],
+#             phi=0.,
+#             alpha=0,
+#             other_beam_q0=1.,
+#             slices_other_beam_Sigma_11=[1],
+#             slices_other_beam_Sigma_12=[0],
+#             slices_other_beam_Sigma_22=[0],
+#             slices_other_beam_Sigma_33=[1],
+#             slices_other_beam_Sigma_34=[0],
+#             slices_other_beam_Sigma_44=[0],
+#         ),
+#         # Doesn't work because fieldmap in different buffer
+#         xf.ElectronCloud(fieldmap=xf.TriCubicInterpolatedFieldMap(
+#             x_range=(0.,1.), nx=10,
+#             y_range=(0.,1.), ny=10,
+#             z_range=(0.,1.), nz=10
+#         )),
+#         # Doesn't work because issue with definition of atomicAdd when generating executable
+#         xf.ElectronLensInterpolated(
+#             x_range=(0.,1.), nx=10,
+#             y_range=(0.,1.), ny=10
+#         ),
+#         # Doesn't work because issue with definition of atomicAdd when generating executable
+#         xf.SpaceCharge3D(
+#             x_range=(0.,1.), nx=10,
+#             y_range=(0.,1.), ny=10,
+#             z_range=(0.,1.), nz=10
+#         ),
         xc.BlackAbsorber(length=1),
         xc.EverestCollimator(length=1, material=xc.materials.Silicon),
-        xc.EverestCrystal(length=1, material=xc.materials.SiliconCrystal)
-        #xf.SpaceChargeBiGaussian( # Not working!
-        #    longitudinal_profile=xf.LongitudinalProfileQGaussian(
-        #        number_of_particles=0, sigma_z=1)
-        #),
+        xc.EverestCrystal(length=1, material=xc.materials.SiliconCrystal),
+        xt.LimitRect(),
+        xt.LimitRacetrack(),
+        xt.LimitEllipse(),
+        xt.LimitPolygon(x_vertices=[0.,1.], y_vertices=[0.,1.]),
+        xt.LimitRectEllipse(),
+        xt.LongitudinalLimitRect()
     ])
 
     _context = xo.ContextCpu()
diff --git a/xboinc/simulation_io/input.py b/xboinc/simulation_io/input.py
@@ -4,12 +4,13 @@
 # ######################################### #
 
 import numpy as np
+from pathlib import Path
 
 import xobjects as xo
 import xpart as xp
 import xtrack as xt
 
-from ..executable import get_default_tracker
+from .default_tracker import get_default_tracker
 from .output import SimState
 
 _default_tracker, _default_config_hash = get_default_tracker()
@@ -24,7 +25,7 @@ class SimConfig(xo.Struct):
     sim_state        = xo.Ref(SimState)
 
     @classmethod
-    def build(name, num_turns, line, particles, checkpoint_every=-1):
+    def build(cls, *, num_turns, line, particles, checkpoint_every=-1, filename=None):
         # Assemble data structure
         simbuf = xo.ContextCpu().new_buffer()
         sim_config = cls(_buffer=simbuf)
@@ -34,18 +35,19 @@ def build(name, num_turns, line, particles, checkpoint_every=-1):
         if _default_config_hash not in line.tracker._tracker_data_cache:
             raise RuntimeError('Tracker data for default config not found')
         tracker_data = line.tracker._tracker_data_cache[_default_config_hash]
-        version, part = get_boinc_particles_from_particles(particles)
-        sim_state = build_simulation_state(version=version, _buffer=simbuf,
-                                         particles=part, i_turn=0)
+        sim_state = SimState(particles=particles, i_turn=0, _buffer=simbuf)
         sim_config.line_metadata = tracker_data._element_ref_data
         sim_config.num_turns = num_turns
         sim_config.num_elements = len(line.element_names)
         sim_config.checkpoint_every = checkpoint_every
         sim_config.sim_state = sim_state._xobject
+        # TODO: can we set size at construction?
         sim_state.size = sim_state._xobject._size # store size of sim_state
         assert sim_config._offset == 0
         assert sim_config._fields[0].offset == 0
         # Write sim buffer to file
-        with open(name, 'wb') as fid:
-            fid.write(simbuf.buffer.tobytes())
+        if filename is not None:
+            filename = Path(filename).expanduser().resolve()
+            with filename.open('wb') as fid:
+                fid.write(simbuf.buffer.tobytes())
         return sim_config
diff --git a/xboinc/simulation_io/output.py b/xboinc/simulation_io/output.py
@@ -35,11 +35,11 @@ class SimVersion(xo.HybridClass):
 
     def __init__(self, **kwargs):
         if '_xobject' not in kwargs:
-            kwargs['xboinc'] = _version_to_int(version)
+            kwargs['xboinc_version'] = _version_to_int(__version__)
         super().__init__(**kwargs)
 
     def assert_version(self):
-        if _version_to_int(__version__) != self.xboinc_version:
+        if int(_version_to_int(__version__)/1000) != int(self.xboinc_version/1000):
             error += f"Incompatible xboinc version! Output file needs "
             error +  f"{_int_to_version(self.xboinc_version)}, "
             error +  f"but current version is {__version__}.\n"
@@ -57,10 +57,11 @@ class SimState(xo.HybridClass):
     def __init__(self, **kwargs):
         if '_xobject' not in kwargs:
             kwargs['version'] = SimVersion()
+        # TODO: can we set size at construction?
         super().__init__(**kwargs)
 
     @classmethod
-    def from_binary(filename, offset=0):
+    def from_binary(cls, filename, offset=0):
         # Read binary
         filename = Path(filename)
         with filename.open('rb') as fid:

Original file line number	Diff line number	Diff line change
`@@ -4,4 +4,3 @@`
`4`	`4`	`# ######################################### #`
`5`	`5`
`6`	`6`	`from .generate_executable_source import generate_executable_source, generate_executable`
`7`		`-from .default_tracker import get_default_tracker`