Merge remote-tracking branch 'origin/feat/scheduling_test' into feat/linspace_measurements/atssimple

Author: maxbeer99

qupulse/examples/expectation_checker_meta.py

@@ -0,0 +1,108 @@
+# -*- coding: utf-8 -*-
+import numpy as np
+
+from qupulse.pulses import PointPT, ConstantPT, RepetitionPT, ForLoopPT, TablePT,\
+    FunctionPT, AtomicMultiChannelPT, SequencePT, MappingPT, ParallelConstantChannelPT,\
+    SchedulerPT
+from qupulse.program.linspace import LinSpaceBuilder, to_increment_commands
+from qupulse.program.multi import MultiProgramBuilder
+from qupulse.program.loop import LoopBuilder
+
+from qupulse.plotting import plot
+from qupulse.utils import to_next_multiple
+
+from expectation_checker_meta import ExpectationChecker, HDAWGAlazar
+
+#%% Get Devices
+
+# Get the Alazar and the HDAWG in hardcoded configuration (2V p-p range)
+# Connect any Marker from the HDAWG to Trig in of the Alazar
+# Connect one dummy channel from HDAWG to the external clock of the Alazar,
+# then set 0.5V-10MHz-oscillator on this channel (e.g. in HDAWG-Webinterface)
+
+ha = HDAWGAlazar("DEVXXXX","USB",)
+
+#%% Example pulse definitions
+
+# PulseTemplates must be defined on channels named as 'ZI0_X', X in A to H
+# and possibly marker channel(s), as denoted +'_MARKER_FRONT'
+# channel subsets {AB,CD,EF,GH} to be referenced as f'expcheck_{i}', i in range(4)
+# (ensure correct mapping to Alazar)
+# Markers need to be added on the according channel to trigger the alazar
+
+        
+class MultiSchedule():
+    def __init__(self,base_time=1e2,rep_factor=3):
+        
+        init_chs = {'ZI0_A','ZI0_B','ZI0_A_MARKER_FRONT','ZI0_B_MARKER_FRONT',}
+        manip_chs = {'ZI0_C','ZI0_D','ZI0_C_MARKER_FRONT','ZI0_D_MARKER_FRONT',}
+        
+        zone1 = {'expcheck_0': init_chs, 'expcheck_1': manip_chs,}
+        
+        init_chs_2 = {'ZI0_E','ZI0_F','ZI0_E_MARKER_FRONT','ZI0_F_MARKER_FRONT',}
+        manip_chs_2 = {'ZI0_G','ZI0_H','ZI0_G_MARKER_FRONT','ZI0_H_MARKER_FRONT',}
+        
+        top_channel_structure = {
+                                'zone1': zone1,
+                                'expcheck_2': init_chs_2,
+                                }
+        
+        markers_ab = {f'ZI0_{x}_MARKER_FRONT':1.0 for x in 'AB'}#|{f'ZI0_{x}_MARKER_BACK':0.0 for x in 'AB'}
+        markers_cd = {f'ZI0_{x}_MARKER_FRONT':1.0 for x in 'CD'}#|{f'ZI0_{x}_MARKER_BACK':0.0 for x in 'CD'}
+        markers_ef = {f'ZI0_{x}_MARKER_FRONT':1.0 for x in 'EF'}#|{f'ZI0_{x}_MARKER_BACK':0.0 for x in 'EF'}
+        markers_ab_to_ef = {ch1:ch2 for ch1,ch2 in zip(markers_ab.keys(),markers_ef.keys())}
+        
+        some_pt = ParallelConstantChannelPT("amp*1/8"*FunctionPT("sin(t/100)","t_sin",channel='a'),{'b':-0.5}|markers_ab)
+        some_pt2 = ParallelConstantChannelPT("amp*1/8"*FunctionPT("sin(t/100)","t_sin",channel='a'),{'b':-0.5}|markers_ab)
+        some_other_pt2 = ParallelConstantChannelPT("amp*1/8"*FunctionPT("sin(t/200)","t_sin",channel='c'),{'d':-0.6}|markers_cd)
+
+        
+        schedule = SchedulerPT(zone1)
+        top_schedule = SchedulerPT(top_channel_structure)
+        
+        s1 = schedule.add_pt(MappingPT(some_pt,channel_mapping=dict(a='ZI0_A',b='ZI0_B')), None)
+        s2 = schedule.add_pt(MappingPT(some_pt2,channel_mapping=dict(a='ZI0_A',b='ZI0_B')),s1)
+        s3 = schedule.add_pt(MappingPT(some_other_pt2,channel_mapping=dict(c='ZI0_C',d='ZI0_D')),s1,rel_time=80+16.)
+        s3 = schedule.add_pt(MappingPT(some_other_pt2,channel_mapping=dict(c='ZI0_C',d='ZI0_D'),parameter_mapping=dict(t_sin="3.5*t_sin",amp='0.7')),s3,rel_time=4*80+16.)
+
+
+        looped = ForLoopPT(schedule, 'amp', 5)
+        # looped = RepetitionPT(schedule,5)
+        
+        n1 = top_schedule.add_pt(looped, None)
+        n2 = top_schedule.add_pt(MappingPT(some_pt2,parameter_mapping={'amp':0.5,'t_sin':1e4},channel_mapping=dict(a='ZI0_E',b='ZI0_F')|markers_ab_to_ef), n1, ref_point="start",rel_time=-128.)
+        # n2 = top_schedule.add_pt(MappingPT(some_pt2,parameter_mapping={'amp':0.5,'t_sin':1e4},channel_mapping=dict(a='ZI0_G',b='ZI0_H')), n1, ref_point="start",rel_time=192.)
+                                             
+        self.pulse_template = MappingPT(top_schedule,parameter_mapping=dict(t_sin=192.))
+        self.channel_subsets = top_channel_structure
+        
+
+#%% Instantiate Checker
+
+# select exemplary pulse
+pulse = MultiSchedule()
+
+# # Define a program builder to test program with:
+default_program_builder = LinSpaceBuilder(
+    #set to True to ensure triggering at Program start if program starts with constant pulse
+    play_marker_when_constant=True,
+    #in case stepped repetitions are needed, insert variables here: 
+    # to_stepping_repeat={'example',},
+    # to_stepping_repeat={'amp','amp2','off','freq'},
+    )
+
+program_builder = MultiProgramBuilder.from_mapping(default_program_builder, channel_subsets=pulse.channel_subsets)
+
+# Data will be saved as xr.Dataset in save_path
+# data_offsets corrects for offsets in Alazar (not in saved data, only in plotting)
+checker = ExpectationChecker(ha, pulse.pulse_template,
+                             program_builder=program_builder,
+                             save_path=r"D:\2302_experiment_env\qupulse_dev_scripts\expcheck",
+                             data_offsets={'t_offset':-100.,'v_offset':0.008,'v_scale':0.9975}
+                             )
+
+assert float(pulse.pulse_template.duration) < 1e7, "Ensure you know what you're doing when recording long data"
+
+#%% Run the checker
+
+checker.run()

qupulse/examples/expectation_checker_meta_example_script.py

@@ -1,11 +1,13 @@
-from typing import NamedTuple, Set, Callable, Dict, Tuple, Union, Iterable, Any, Mapping
+from typing import NamedTuple, Set, Callable, Dict, Tuple, Union, Iterable, Any, Mapping, Self
 from collections import defaultdict
+from dataclasses import dataclass
 import warnings
 import numbers
 
 from qupulse.hardware.awgs.base import AWG
 from qupulse.hardware.dacs import DAC
 from qupulse.program.loop import Loop
+from qupulse.program.multi import MultiProgram
 
 from qupulse.utils.types import ChannelID
 
@@ -313,8 +315,171 @@ def registered_programs(self) -> Dict[str, RegisteredProgram]:
         return self._registered_programs
 
 
+@dataclass
+class RegisteredMultiProgram:
+    program: MultiProgram
+    measurement_windows: Dict[str, Tuple[np.ndarray, np.ndarray]]
+    run_callback: Callable
+    hw_setups_to_utilize: Set[HardwareSetup]
+    dacs_to_arm: Set[DAC]
 
 
+class MetaHardwareSetup:
+
+    def __init__(self):
+        # self._channel_map: Dict[str, Set[ChannelID]] = dict()
+        self._setup_map: Dict[str, HardwareSetup] = dict()
+        
+
+        self._registered_programs: Dict[str, RegisteredMultiProgram] = dict()
+        
+        self._alibi_measurement_setup = HardwareSetup()
+        
+    def register_program(self,
+                         name: str,
+                         program: MultiProgram,
+                         run_callback: Callable = lambda: None,
+                         update: bool = False,
+                         flatten_structure: bool = False,
+                         measurements: Mapping[str, Tuple[np.ndarray, np.ndarray]] = None
+                         ) -> None:
+        
+        hw_setups_to_utilize = set()
+        
+        if flatten_structure:
+            program_map = program._flattened_program_map
+        else:
+            program_map = program.program_map
+        
+        ############
+        #GET MEASUREMENT WINDOWS, collect from all subprograms (can have different ones...)
+        # measurements: Mapping[str, Tuple[np.ndarray, np.ndarray]]
+        if measurements is None:
+            measurements = program.get_measurement_windows(drop=True)
+        
+        self._alibi_measurement_setup.register_program(name, None,
+                                                       run_callback=lambda: None,
+                                                       channels=set(),measurements=measurements
+                                                       )
+        
+        #END MEAS.        
+        ############
+        
+        
+        
+        
+        for s_ident,prog in program_map.items():
+            self.setup_map[s_ident].register_program(s_ident+'_'+name, prog,
+                                                     update=update,run_callback=lambda:None,
+                                                     measurements={})
+            hw_setups_to_utilize.add(self.setup_map[s_ident])
+    
+        self._registered_programs[name] = RegisteredMultiProgram(program,
+                                                                 self._alibi_measurement_setup.registered_programs[name].measurement_windows, #meas todo TODOTODO exract from alibi setup
+                                                                 run_callback,
+                                                                 hw_setups_to_utilize,
+                                                                 dacs_to_arm=self._alibi_measurement_setup.registered_programs[name].dacs_to_arm,
+                                                                 # dacs_to_arm=set(affected_dacs.keys())
+                                                                 )
+        
+        return
+    
+    @property
+    def _measurement_map(self) -> Dict[str, Set[MeasurementMask]]:
+        return self._alibi_measurement_setup._measurement_map
+
+    
+    def set_measurement(self, measurement_name: str,
+                        measurement_mask: Union[MeasurementMask, Iterable[MeasurementMask]],
+                        allow_multiple_registration: bool=False):
+        return self._alibi_measurement_setup.set_measurement(measurement_name, measurement_mask, allow_multiple_registration)
+    
+    
+    def remove_program(self, name: str):
+        if name in self.registered_programs:
+            for setup in self.setup_map.values():
+                if name in setup.registered_programs:
+                    setup.remove_program(name)
+            self._alibi_measurement_setup.remove_program(name)
+                
+        
+    def clear_programs(self):
+        for setup in self._setup_map.values():
+            setup.clear_programs()
+        
+        self._alibi_measurement_setup.clear_programs()
+        self._registered_programs = {}
+
+    @property
+    def known_hw_setups(self) -> Set[Union[HardwareSetup,"MetaHardwareSetup"]]:
+        return set(self._setup_map.values())
+    
+    @property
+    def setup_map(self) -> Dict[str,Union[HardwareSetup,"MetaHardwareSetup"]]:
+        return self._setup_map
+    
+    @property
+    def known_awgs(self) -> Set[AWG]:
+        return set().union(*[s.known_awgs for s in self._setup_map.values()])
+    
+    @property
+    def known_dacs(self) -> Set[DAC]:
+        return self._alibi_measurement_setup.known_dacs
+    
+    def arm_program(self, name: str) -> None:
+        """Assert program is in memory. Hardware will wait for trigger event"""
+        if name not in self._registered_programs:
+            raise KeyError('{} is not a registered program'.format(name))
+
+        hw_setups_to_utilize = self._registered_programs[name].hw_setups_to_utilize
+        for key,setup in self._setup_map.items():
+            if setup in hw_setups_to_utilize:
+                setup.arm_program(key+'_'+name)
+            else:
+                # The other AWGs should ignore the trigger
+                # awg.arm(None)
+                pass
+        self._alibi_measurement_setup.arm_program(name)
+        # for dac in dacs_to_arm:
+        #     dac.arm_program(name)
+    @property 
+    def _channel_map(self) -> Dict[str, Set[ChannelID]]:
+        return {k: s.registered_channels.keys() for k,s in self._setup_map.items()}
+        
+    def run_program(self, name) -> None:
+        """Calls arm program and starts it using the run callback"""
+        self.arm_program(name)
+        self._registered_programs[name].run_callback()
+
+    def add_setup(self, identifier: str,
+                    setup: Union[HardwareSetup,"MetaHardwareSetup"],
+                    allow_multiple_registration: bool=False) -> None:
+        
+        assert identifier not in self._setup_map
+        assert setup not in self._setup_map.values()
+        
+        assert not any(s in self.known_awgs for s in setup.known_awgs)
+        
+        self._setup_map[identifier] = setup
+
+    def rm_setup(self, identifier: str) -> None:
+        self._setup_map.pop(identifier)
+        
+    @property
+    def all_registered_channels(self) -> Dict[ChannelID, Set[_SingleChannel]]:
+        return {k:v for setup in self._setup_map.values() for k,v in setup.registered_channels().items()}
+
+    # def update_parameters(self, name: str, parameters: Mapping[str, numbers.Real]):
+    #     *_, awgs, dacs = self._registered_programs[name]
+
+    #     for awg in self.known_awgs:
+    #         if awg in awgs:
+    #             awg.set_volatile_parameters(name, parameters)
+
+    @property
+    def registered_programs(self) -> Dict[str, RegisteredMultiProgram]:
+        return self._registered_programs
+
 
 
 

qupulse/hardware/setup.py

@@ -74,7 +74,9 @@ def with_iteration(self, index_name: str, rng: range,
     def evaluate_nested_stepping(self, scope: Scope, parameter_names: set[str]) -> bool:
         return False
 
-    def to_program(self, defined_channels: Set[ChannelID]) -> Optional[Program]:
+    def to_program(self,
+                   # defined_channels: Set[ChannelID]
+                   ) -> Optional[Program]:
         """Further addition of new elements might fail after finalizing the program."""