match chroma in 4D mode

toolkit/cpymad_utils.py

@@ -20,6 +20,12 @@ def set_up_lattice(mad_instance, reference_parameter, seq_path):
 def match_chromaticity(mad_instance, chromas, sextupole_script_path):
     mad_instance.call(sextupole_script_path)
     mad_instance.input(
+"""VOLTCA1SAVE = VOLTCA1; 
+VOLTCA2SAVE = VOLTCA2; 
+
+VOLTCA1 = 0;
+VOLTCA2 = 0;""")
+    mad_instance.input(
 f"""MATCH, SEQUENCE=FCCEE_P_RING, CHROM;
 VARY, NAME=KN_SF;
 VARY, NAME=KN_SD;
@@ -30,6 +36,10 @@ GLOBAL, DQ2={chromas['QY']};
 LMDIF, CALLS=10000;
 ENDMATCH;
 """)
+    
+    mad_instance.input(
+"""VOLTCA1 = VOLTCA1SAVE;
+VOLTCA2 = VOLTCA2SAVE;""")
 
 
 def taper_lattice(mad_instance):