add natoms from builder (#451)

* add natoms from builder

* rename natom -> n_atoms

---------

Co-authored-by: Kai-Hsin Wu <khwu@KHWus-MacBook-Pro.local>
Co-authored-by: Phillip Weinberg <weinbe58@gmail.com>

src/bloqade/builder/compile/trait.py

@@ -35,4 +35,14 @@ def parse_program(self) -> ir.Program:
 
 
 class Parse(ParseRegister, ParseSequence, ParseProgram):
-    pass
+    @property
+    def n_atoms(self):
+        from .builder import Parser
+
+        ps = Parser(self)
+        ps.read_register()
+
+        if isinstance(ps.register, ir.location.ParallelRegister):
+            return ps.register._register.n_atoms
+        else:
+            return ps.register.n_atoms

tests/test_program_visitor.py

@@ -50,3 +50,29 @@ def test_base_program_visitor():
     scaled_loc = list(seq2.pulses[rydberg].fields[Detuning()].value.keys())[0]
     with pytest.raises(NotImplementedError):
         pvis.visit(scaled_loc)
+
+
+def test_get_natoms():
+    prog = (
+        Square(5)
+        .rydberg.detuning.uniform.piecewise_linear(
+            [0.1, 0.9, 0.1],
+            [
+                -10,
+                -10,
+                10,
+                10,
+            ],
+        )
+        .hyperfine.detuning.uniform.piecewise_linear(
+            [0.1, 0.9, 0.1],
+            [
+                -10,
+                -10,
+                10,
+                10,
+            ],
+        )
+    )
+
+    assert prog.n_atoms == 25