Beam blanked unblanked

docs/tem_api.md

@@ -183,6 +183,16 @@ To unblank the beam:
 ```python
 ctrl.beam.unblank()
 ```
+To blank the beam temporarily:
+```python
+with ctrl.beam.blanked():
+    ...
+```
+To unblank the beam temporarily:
+```python
+with ctrl.beam.unblanked():
+    ...
+```
 To get the state of the beam blanker:
 ```python
 status = ctrl.beam.status          # str

src/instamatic/controller.py

@@ -651,17 +651,10 @@ def get_image(
 
         h = self.to_dict(*header_keys) if header_keys else {}
 
-        if self.autoblank:
-            self.beam.unblank()
-
-        h['ImageGetTimeStart'] = time.perf_counter()
-
-        arr = self.get_rotated_image(exposure=exposure, binsize=binsize)
-
-        h['ImageGetTimeEnd'] = time.perf_counter()
-
-        if self.autoblank:
-            self.beam.blank()
+        with self.beam.unblanked(condition=self.autoblank):
+            h['ImageGetTimeStart'] = time.perf_counter()
+            arr = self.get_rotated_image(exposure=exposure, binsize=binsize)
+            h['ImageGetTimeEnd'] = time.perf_counter()
 
         h['ImageGetTime'] = time.time()
         h['ImageExposureTime'] = exposure
@@ -757,11 +750,7 @@ def get_movie(
         header_common['ImageCameraDimensions'] = self.cam.get_camera_dimensions()
 
         gen = self.cam.get_movie(n_frames=n_frames, exposure=exposure, binsize=binsize)
-
-        try:
-            if self.autoblank:
-                self.beam.unblank()
-
+        with self.beam.unblanked(condition=self.autoblank):
             for _ in range(n_frames):
                 # The generator `gen` starts collecting only when the first `next` is called.
                 # Request the next image, expect it in the future, get header in the meantime
@@ -786,11 +775,7 @@ def get_movie(
                 rotate_image(img, mode=mode, mag=mag)
                 header['ImageResolution'] = img.shape
                 yield img, header
-
-        finally:
-            gen.close()
-            if self.autoblank:
-                self.beam.blank()
+        gen.close()
 
     def store_diff_beam(self, name: str = 'beam', save_to_file: bool = False):
         """Record alignment for current diffraction beam. Stores Guntilt (for

src/instamatic/experiments/cred_gatan/experiment.py

@@ -183,15 +183,13 @@ def start_collection(
             # Center crystal position
             if self.mode == 'diff':
                 self.ctrl.difffocus.defocus(self.defocus_offset)
-            self.ctrl.beam.unblank()
-
-            input('Move SAED aperture to crystal and press <ENTER> to measure! ')
+            with self.ctrl.beam.unblanked():
+                input('Move SAED aperture to crystal and press <ENTER> to measure! ')
 
             # cannot do this while lieview is running
             # img1 = self.ctrl.get_raw_image()
             # write_tiff(self.path / "image_before.tiff", img1)
 
-            self.ctrl.beam.blank()
             if self.mode == 'diff':
                 self.ctrl.difffocus.refocus()
             time.sleep(3)
@@ -305,12 +303,9 @@ def start_collection(
             # Center crystal position
             if self.mode == 'diff':
                 self.ctrl.difffocus.defocus(self.defocus_offset)
-            self.ctrl.beam.unblank()
-
-            img2 = self.ctrl.get_rotated_image()
-            write_tiff(self.path / 'image_after.tiff', img2)
-
-            self.ctrl.beam.blank()
+            with self.ctrl.beam.unblanked():
+                img2 = self.ctrl.get_rotated_image()
+                write_tiff(self.path / 'image_after.tiff', img2)
             if self.mode == 'diff':
                 self.ctrl.difffocus.refocus()
 

src/instamatic/experiments/cred_tvips/experiment.py

@@ -472,12 +472,9 @@ def start_collection(
             # Center crystal position
             if self.mode == 'diff':
                 self.ctrl.difffocus.defocus(self.defocus_offset)
-            self.ctrl.beam.unblank()
-
-            img2 = self.ctrl.get_rotated_image()
-            write_tiff(self.path / 'image_after.tiff', img2)
-
-            self.ctrl.beam.blank()
+            with self.ctrl.beam.unblanked():
+                img2 = self.ctrl.get_rotated_image()
+                write_tiff(self.path / 'image_after.tiff', img2)
             if self.mode == 'diff':
                 self.ctrl.difffocus.refocus()
 

src/instamatic/microscope/components/states.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import time
+from contextlib import contextmanager
+from typing import Any, ContextManager
 
 from instamatic.microscope.base import MicroscopeBase
 
@@ -62,6 +64,30 @@ def unblank(self, delay: float = 0.0) -> None:
         if delay:
             time.sleep(delay)
 
+    @contextmanager
+    def blanked(self, condition=True, delay: float = 0.0) -> ContextManager[None]:
+        """Temporarily blank the beam using a `with blanked` statement."""
+        was_blanked_before = self.is_blanked
+        try:
+            if condition and not was_blanked_before:
+                self.blank(delay=delay)
+            yield
+        finally:
+            if condition and not was_blanked_before:
+                self.unblank(delay=delay)
+
+    @contextmanager
+    def unblanked(self, condition=True, delay: float = 0.0) -> ContextManager[None]:
+        """Temporarily unblank the beam using a `with unblanked` statement."""
+        was_blanked_before = self.is_blanked
+        try:
+            if condition and was_blanked_before:
+                self.unblank(delay=delay)
+            yield
+        finally:
+            if condition and was_blanked_before:
+                self.blank(delay=delay)
+
     def set(self, state: str, delay: float = 0.0):
         """Set state of the beam, with optional delay in ms."""
         index = self._states.index(state)

src/instamatic/processing/flatfield.py

@@ -134,26 +134,23 @@ def collect_flatfield(
     np.save(fp, deadpixels)
 
     if collect_darkfield:
-        ctrl.beam.blank()
-
-        buffer = []
-
-        print('\nCollecting darkfield images')
-        for n in tqdm(range(frames)):
-            outfile = drc / f'darkfield_{n:04d}.tiff' if save_images else None
-            img, h = ctrl.get_image(
-                exposure=exposure,
-                binsize=binsize,
-                out=outfile,
-                comment=f'Dark field #{n:04d}',
-                header_keys=None,
-            )
-            buffer.append(img)
-
-        d = np.mean(buffer, axis=0)
-        d = remove_deadpixels(d, deadpixels=deadpixels)
-
-        ctrl.beam.unblank()
+        with ctrl.beam.blanked():
+            buffer = []
+
+            print('\nCollecting darkfield images')
+            for n in tqdm(range(frames)):
+                outfile = drc / f'darkfield_{n:04d}.tiff' if save_images else None
+                img, h = ctrl.get_image(
+                    exposure=exposure,
+                    binsize=binsize,
+                    out=outfile,
+                    comment=f'Dark field #{n:04d}',
+                    header_keys=None,
+                )
+                buffer.append(img)
+
+            d = np.mean(buffer, axis=0)
+            d = remove_deadpixels(d, deadpixels=deadpixels)
 
         fd = drc / f'darkfield_{ctrl.cam.name}_{date}.tiff'
         write_tiff(fd, d, header={'deadpixels': deadpixels})

tests/test_ctrl.py

@@ -183,6 +183,12 @@ def test_beam(ctrl):
     beam.blank()
     assert beam.is_blanked
 
+    with beam.blanked():
+        assert beam.is_blanked
+
+    with beam.unblanked():
+        assert not beam.is_blanked
+
     beam.set(unblanked)
     assert beam.state == unblanked