Atom arrangement

def add_position(
    self,
    position: Union[
        PositionArray,
        List[Tuple[ScalarType, ScalarType]],
        Tuple[ScalarType, ScalarType],
    ],
    filling: Optional[Union[BoolArray, List[bool], bool]] = None,
) -> "ListOfLocations":
    """
    Add a position or multiple positions to a pre-existing geometry.

    `add_position` is capable of accepting:
    - A single tuple for one atom coordinate: `(1.0, 2.5)`
    - A list of tuples: `[(0.0, 1.0), (2.0,1.5), etc.]
    - A numpy array of shape (N, 2) where N is the number of atoms

    You may also intersperse variables anywhere a value may be present.

    You can also pass in an optional argument which determines the atom "filling"
    (whether or not at a specified coordinate an atom should be present).

    ### Usage Example:
    ```
    # single coordinate
    >>> reg = start.add_position((0,0))
    # you may chain add_position calls
    >>> reg_plus_two = reg.add_position([(2,2),(5.0, 2.1)])
    # you can add variables anywhere a value may be present
    >>> reg_with_var = reg_plus_two.add_position(("x", "y"))
    # and specify your atom fillings
    >>> reg_with_filling = reg_with_var.add_position([(3.1, 0.0), (4.1, 2.2)],
    [True, False])
    # alternatively you could use one boolean to specify
    # all coordinates should be empty/filled
    >>> reg_with_more_filling = reg_with_filling.add_positions([(3.1, 2.9),
    (5.2, 2.2)], False)
    ```

    - Next possible steps are:
    - Continuing to build your geometry via:
        - `...add_position(positions).add_position(positions)`:
            to add more positions
        - `...add_position(positions).apply_defect_count(n_defects)`:
        to randomly drop out n_atoms
        - `...add_position(positions).apply_defect_density(defect_probability)`:
        to drop out atoms with a certain probability
        - `...add_position(positions).scale(scale)`: to scale the geometry
    - Targeting a level coupling once you're done with the atom geometry:
        - `...add_position(positions).rydberg`: to specify Rydberg coupling
        - `...add_position(positions).hyperfine`: to specify Hyperfine coupling
    - Visualizing your atom geometry:
        - `...add_position(positions).show()`:
        shows your geometry in your web browser

    """

    if is_bearable(position, PositionArray) and is_bearable(
        filling, Optional[BoolArray]
    ):
        return self.add_position_ndarray(position, filling)
    elif is_bearable(position, List[Tuple[ScalarType, ScalarType]]) and is_bearable(
        filling, Optional[List[bool]]
    ):
        return self.add_position_list_tuples(position, filling)
    elif is_bearable(position, Tuple[ScalarType, ScalarType]) and is_bearable(
        filling, Optional[bool]
    ):
        return self.add_position_single_tupe(position, filling)
    else:
        raise TypeError("Invalid input types for add_position provided!")

@beartype
def apply_defect_count(
    self, n_defects: int, rng: np.random.Generator = np.random.default_rng()
):
    """
    Drop `n_defects` atoms from the geometry randomly. Internally this occurs
    by setting certain sites to have a SiteFilling set to false indicating
    no atom is present at the coordinate.

    A default numpy-based Random Number Generator is used but you can
    explicitly override this by passing in your own.

    ### Usage Example:

    ```
    >>> from bloqade.atom_arrangement import Chain
    >>> import numpy as np
    # set a custom seed for a numpy-based RNG
    >>> custom_rng = np.random.default_rng(888)
    # randomly remove two atoms from the geometry
    >>> reg = Chain(11).apply_defect_count(2, custom_rng)
    # you may also chain apply_defect_count calls
    >>> reg.apply_defect_count(2, custom_rng)
    # you can also use apply_defect_count on custom geometries
    >>> from bloqade import start
    >>> start.add_position([(0,0), (1,1)]).apply_defect_count(1, custom_rng)
    ```

    - Next possible steps are:
    - Continuing to build your geometry via:
        - `...apply_defect_count(defect_counts).add_position(positions)`:
            to add more positions
        - `...apply_defect_count(defect_counts)
            .apply_defect_count(n_defects)`: to randomly drop out n_atoms
        - `...apply_defect_count(defect_counts)
            .apply_defect_density(defect_probability)`:
            to drop out atoms with a certain probability
        - `...apply_defect_count(defect_counts).scale(scale)`:
            to scale the geometry
    - Targeting a level coupling once you're done with the atom geometry:
        - `...apply_defect_count(defect_counts).rydberg`: to specify
            Rydberg coupling
        - `...apply_defect_count(defect_counts).hyperfine`:
            to specify Hyperfine coupling
    - Visualizing your atom geometry:
        - `...apply_defect_count(defect_counts).show()`:
            shows your geometry in your web browser
    """

    location_list = []
    for location_info in self.enumerate():
        location_list.append(location_info)

    filled_sites = []

    for index, location_info in enumerate(location_list):
        if location_info.filling is SiteFilling.filled:
            filled_sites.append(index)

    if n_defects >= len(filled_sites):
        raise ValueError(
            f"n_defects {n_defects} must be less than the number of filled sites "
            f"({len(filled_sites)})"
        )

    for _ in range(n_defects):
        index = rng.choice(filled_sites)
        location_list[index] = LocationInfo.create(
            location_list[index].position,
            (False if location_list[index].filling is SiteFilling.filled else True),
        )
        filled_sites.remove(index)

    return ListOfLocations(location_list)

@beartype
def apply_defect_density(
    self,
    defect_probability: float,
    rng: np.random.Generator = np.random.default_rng(),
):
    """
    Drop atoms randomly with `defect_probability` probability (range of 0 to 1).
    Internally this occurs by setting certain sites to have a SiteFilling
    set to false indicating no atom is present at the coordinate.

    A default numpy-based Random Number Generator is used but you can
    explicitly override this by passing in your own.

    ### Usage Example:

    ```
    >>> from bloqade.atom_arrangement import Chain
    >>> import numpy as np
    # set a custom seed for a numpy-based RNG
    >>> custom_rng = np.random.default_rng(888)
    # randomly remove two atoms from the geometry
    >>> reg = Chain(11).apply_defect_density(0.2, custom_rng)
    # you may also chain apply_defect_density calls
    >>> reg.apply_defect_count(0.1, custom_rng)
    # you can also use apply_defect_density on custom geometries
    >>> from bloqade import start
    >>> start.add_position([(0,0), (1,1)])
    .apply_defect_density(0.5, custom_rng)
    ```

    - Next possible steps are:
    - Continuing to build your geometry via:
        - `...apply_defect_count(defect_counts).add_position(positions)`:
        to add more positions
        - `...apply_defect_count(defect_counts).apply_defect_count(n_defects)`:
        to randomly drop out n_atoms
        - `...apply_defect_count(defect_counts)
        .apply_defect_density(defect_probability)`:
        to drop out atoms with a certain probability
        - `...apply_defect_count(defect_counts).scale(scale)`:
        to scale the geometry
    - Targeting a level coupling once you're done with the atom geometry:
        - `...apply_defect_count(defect_counts).rydberg`:
        to specify Rydberg coupling
        - `...apply_defect_count(defect_counts).hyperfine`:
        to specify Hyperfine coupling
    - Visualizing your atom geometry:
        - `...apply_defect_count(defect_counts).show()`:
        shows your geometry in your web browser
    """

    p = min(1, max(0, defect_probability))
    location_list = []

    for location_info in self.enumerate():
        if rng.random() < p:
            location_list.append(
                LocationInfo.create(
                    location_info.position,
                    (
                        False
                        if location_info.filling is SiteFilling.filled
                        else True
                    ),
                )
            )
        else:
            location_list.append(location_info)

    return ListOfLocations(location_list=location_list)

def enumerate(self) -> Generator[LocationInfo, None, None]:
    """enumerate all locations in the register."""
    raise NotImplementedError

def figure(self, fig_kwargs=None, **assignments):
    """obtain a figure object from the atom arrangement."""
    return get_atom_arrangement_figure(self, fig_kwargs=fig_kwargs, **assignments)

def rydberg_interaction(self, **assignments) -> NDArray:
    """calculate the Rydberg interaction matrix.

    Args:
        **assignments: the values to assign to the variables in the register.

    Returns:
        NDArray: the Rydberg interaction matrix in the lower triangular form.

    """

    from bloqade.constants import RB_C6

    # calculate the Interaction matrix
    V_ij = np.zeros((self.n_sites, self.n_sites))
    for i, site_i in enumerate(self.enumerate()):
        pos_i = np.array([float(ele(**assignments)) for ele in site_i.position])

        for j, site_j in enumerate(self.enumerate()):
            if j >= i:
                break  # enforce lower triangular form

            pos_j = np.array([float(ele(**assignments)) for ele in site_j.position])
            r_ij = np.linalg.norm(pos_i - pos_j)

            V_ij[i, j] = RB_C6 / r_ij**6

    return V_ij

@beartype
def scale(self, scale: ScalarType):
    """
    Scale the geometry of your atoms.

    ### Usage Example:
    ```
    >>> reg = start.add_position([(0,0), (1,1)])
    # atom positions are now (0,0), (2,2)
    >>> new_reg = reg.scale(2)
    # you may also use scale on pre-defined geometries
    >>> from bloqade.atom_arrangement import Chain
    # atoms in the chain will now be 2 um apart versus
    # the default 1 um
    >>> Chain(11).scale(2)
    ```

    - Next possible steps are:
    - Continuing to build your geometry via:
        - `...add_position(positions).add_position(positions)`:
            to add more positions
        - `...add_position(positions).apply_defect_count(n_defects)`:
        to randomly drop out n_atoms
        - `...add_position(positions).apply_defect_density(defect_probability)`:
        to drop out atoms with a certain probability
        - `...add_position(positions).scale(scale)`: to scale the geometry
    - Targeting a level coupling once you're done with the atom geometry:
        - `...add_position(positions).rydberg`:
        to specify Rydberg coupling
        - `...add_position(positions).hyperfine`:
        to specify Hyperfine coupling
    - Visualizing your atom geometry:
        - `...add_position(positions).show()`:
        shows your geometry in your web browser

    """

    scale = cast(scale)
    location_list = []
    for location_info in self.enumerate():
        x, y = location_info.position
        new_position = (scale * x, scale * y)
        location_list.append(
            LocationInfo.create(new_position, bool(location_info.filling.value))
        )

    return ListOfLocations(location_list)