예제 #1
0
def test_slice(tmpdir_factory):

    directory = tmpdir_factory.mktemp("proc")
    filename = os.path.join(directory, "temp.h5")
    box = (400, 400, 400)
    centre = (200, 200, 200)
    shape = {"type": "cube", "cube": {"length": 400}}
    sample = parakeet.sample.new(filename, box, centre, shape)
    parakeet.sample.add_single_molecule(sample, "4v5d")

    # Create the system configuration
    system_conf = multem.SystemConfiguration()
    system_conf.precision = "float"
    if multem.is_gpu_available():
        system_conf.device = "device"
    else:
        system_conf.device = "host"

    # Create the input multislice configuration
    n_phonons = 50
    input_multislice = create_input_multislice(n_phonons, False)

    input_multislice.spec_atoms = sample.get_atoms().to_multem()
    input_multislice.spec_lx = sample.containing_box[1][0]
    input_multislice.spec_ly = sample.containing_box[1][1]
    input_multislice.spec_lz = sample.containing_box[1][2]
    input_multislice.spec_dz = 3

    print("Standard")
    output = multem.simulate(system_conf, input_multislice)

    print("Subslicing")

    # Create the input multislice configuration
    n_slices = 4

    # Slice the sample
    subslices = list(
        multem.slice_spec_atoms(
            input_multislice.spec_atoms, input_multislice.spec_lz, n_slices
        )
    )

    # Do the slices simulation
    sliced_output = multem.simulate(system_conf, input_multislice, iter(subslices))

    # Print the difference
    a = np.abs(np.array(output.data[-1].psi_coh)) ** 2
    b = np.abs(np.array(sliced_output.data[-1].psi_coh)) ** 2
    diff = np.max(np.abs(a - b))
예제 #2
0
input_multislice.obj_lens_c_23 = 0.0
input_multislice.obj_lens_phi_23 = 0.0
input_multislice.obj_lens_inner_aper_ang = 0.0
input_multislice.obj_lens_outer_aper_ang = 0.0

# defocus spread function
dsf_sigma = multem.iehwgd_to_sigma(32)
input_multislice.obj_lens_dsf_sigma = dsf_sigma
input_multislice.obj_lens_dsf_npoints = 5

# zero defocus reference
input_multislice.obj_lens_zero_defocus_type = "First"
input_multislice.obj_lens_zero_defocus_plane = 0

# Do the simulation
output_multislice = multem.simulate(system_conf, input_multislice)

data = {
    "dx":
    output_multislice.dx,
    "dy":
    output_multislice.dy,
    "x":
    numpy.array(output_multislice.x, dtype=numpy.float64),
    "y":
    numpy.array(output_multislice.y, dtype=numpy.float64),
    "thick":
    numpy.array(output_multislice.thick, dtype=numpy.float64),
    "data": [{
        "m2psi_tot": numpy.array(d.m2psi_tot, numpy.float64)
    } for d in output_multislice.data],
def compute_exit_wave(atom_data, pixel_size):
    """
    Compute the exit wave

    """

    # Get the dimensions
    x_min = atom_data.data["x"].min()
    x_max = atom_data.data["x"].max()
    y_min = atom_data.data["y"].min()
    y_max = atom_data.data["y"].max()
    z_min = atom_data.data["z"].min()
    z_max = atom_data.data["z"].max()
    x_size = x_max - x_min
    y_size = y_max - y_min
    select = ((atom_data.data["x"] > x_min + x_size / 6)
              & (atom_data.data["x"] < x_max - x_size / 6)
              & (atom_data.data["y"] > y_min + y_size / 6)
              & (atom_data.data["y"] < y_max - y_size / 6))
    atom_data = parakeet.sample.AtomData(data=atom_data.data[select])
    x_min = atom_data.data["x"].min()
    x_max = atom_data.data["x"].max()
    y_min = atom_data.data["y"].min()
    y_max = atom_data.data["y"].max()
    z_min = atom_data.data["z"].min()
    z_max = atom_data.data["z"].max()
    x_size = x_max - x_min
    y_size = y_max - y_min
    z_size = z_max - z_min

    # Translate to centre
    x_box_size = x_size
    y_box_size = y_size
    z_box_size = z_size

    # Create the system configuration
    system_conf = multem.SystemConfiguration()
    system_conf.precision = "float"
    system_conf.device = "device"

    # Create the input multislice configuration
    input_multislice = create_input_multislice()

    # Compute the number of pixels
    nx = next_power_2(int(x_box_size / pixel_size))
    ny = next_power_2(int(y_box_size / pixel_size))
    assert nx <= 4096
    assert ny <= 4096
    x_box_size = nx * pixel_size
    y_box_size = ny * pixel_size
    x_trans = (x_box_size - x_size) / 2.0 - x_min
    y_trans = (y_box_size - y_size) / 2.0 - y_min
    z_trans = (z_box_size - z_size) / 2.0 - z_min
    atom_data.translate((x_trans, y_trans, z_trans))

    # Create the specimen size
    input_multislice.nx = nx
    input_multislice.ny = ny
    input_multislice.spec_lx = x_box_size
    input_multislice.spec_ly = y_box_size
    input_multislice.spec_lz = z_box_size
    input_multislice.spec_dz = 5

    # Set the specimen atoms
    input_multislice.spec_atoms = atom_data.to_multem()

    # Run the simulation
    output_multislice = multem.simulate(system_conf, input_multislice)

    # Get the image
    physical_image = numpy.array(output_multislice.data[0].psi_coh).T

    # Create the masker
    masker = multem.Masker(input_multislice.nx, input_multislice.ny,
                           pixel_size)

    # Create the size of the cuboid
    masker.set_cuboid(
        (
            x_box_size / 2 - x_size / 2,
            y_box_size / 2 - y_size / 2,
            z_box_size / 2 - z_size / 2,
        ),
        (x_size, y_size, z_size),
    )

    # Run the simulation
    input_multislice.spec_atoms = multem.AtomList()
    output_multislice = multem.simulate(system_conf, input_multislice, masker)

    # Get the image
    random_image = numpy.array(output_multislice.data[0].psi_coh).T

    # Return the images
    x0 = numpy.array(
        (x_box_size / 2 - x_size / 2, y_box_size / 2 - y_size / 2))
    x1 = numpy.array(
        (x_box_size / 2 + x_size / 2, y_box_size / 2 + y_size / 2))
    return physical_image, random_image, x0, x1
예제 #4
0
import multem

input_multislice = multem.Input()
config = multem.SystemConfiguration()

# print(multem.is_gpu_available())
# if not multem.is_gpu_available():
config.device = "host"
result = multem.simulate(config, input_multislice)

# print(result)
    rms3d = 0.085

    (
        input_multislice.spec_atoms,
        input_multislice.spec_lx,
        input_multislice.spec_ly,
        input_multislice.spec_lz,
        a,
        b,
        c,
        input_multislice.spec_dz,
    ) = SrTiO3001_crystal(na, nb, nc, ncu, rms3d)

    print("Standard")
    start_time = time.time()
    ewrs = multem.simulate(system_conf, input_multislice)
    print("Time taken: ", time.time() - start_time)

    print("Subslicing")
    start_time = time.time()

    # Create the input multislice configuration
    n_slices = 4

    subslices = iter(
        list(
            multem.slice_spec_atoms(input_multislice.spec_atoms,
                                    input_multislice.spec_lz, n_slices)))
    output_multislice = multem.simulate(system_conf, input_multislice,
                                        subslices)
    def __call__(self, index):
        """
        Simulate a single frame

        Args:
            simulation (object): The simulation object
            index (int): The frame number

        Returns:
            tuple: (angle, image)

        """

        # Get the specimen atoms
        logger.info(f"Simulating image {index+1}")

        # Get the rotation angle
        angle = self.scan.angles[index]
        position = self.scan.positions[index]

        # Add the beam drift
        if (
            self.microscope.beam.drift
            and not self.microscope.beam.drift["type"] is None
        ):
            if self.microscope.beam.drift["type"] == "random":
                shiftx, shifty = numpy.random.normal(
                    0, self.microscope.beam.drift["magnitude"], size=2
                )
                logger.info("Adding drift of %f, %f " % (shiftx, shifty))
            elif self.microscope.beam.drift["type"] == "sinusoidal":
                shiftx = sin(angle * pi / 180) * self.microscope.beam.drift["magnitude"]
                shifty = shiftx
                logger.info("Adding drift of %f, %f " % (shiftx, shifty))
            else:
                raise RuntimeError("Unknown drift type")
        else:
            shiftx = 0
            shifty = 0

        # The field of view
        nx = self.microscope.detector.nx
        ny = self.microscope.detector.ny
        pixel_size = self.microscope.detector.pixel_size
        origin = numpy.array(self.microscope.detector.origin)
        margin = self.simulation["margin"]
        padding = self.simulation["padding"]
        x_fov = nx * pixel_size
        y_fov = ny * pixel_size
        margin_offset = margin * pixel_size
        # padding_offset = padding * pixel_size
        offset = (padding + margin) * pixel_size

        # Set the rotation angle
        # input_multislice.spec_rot_theta = angle
        # input_multislice.spec_rot_u0 = simulation.scan.axis

        # Create the sample extractor
        x0 = numpy.array((-margin_offset, position - margin_offset))
        x1 = numpy.array((x_fov + margin_offset, position + y_fov + margin_offset))
        x0 += origin
        x1 += origin
        # thickness = self.simulation["division_thickness"]
        # extractor = parakeet.sample.AtomSliceExtractor(
        #    sample=self.sample,
        #    translation=position,
        #    rotation=angle,
        #    x0=x0,
        #    x1=x1,
        #    thickness=thickness,
        # )

        # Create the multem system configuration
        system_conf = create_system_configuration(self.device)

        # The Z centre
        z_centre = self.sample.centre[2]

        # Create the multem input multislice object
        input_multislice = create_input_multislice(
            self.microscope,
            self.simulation["slice_thickness"],
            self.simulation["margin"] + self.simulation["padding"],
            "EWRS",
            z_centre,
        )

        # Set the specimen size
        input_multislice.spec_lx = x_fov + offset * 2
        input_multislice.spec_ly = y_fov + offset * 2
        input_multislice.spec_lz = self.sample.containing_box[1][2]

        # Compute the B factor
        if self.simulation["radiation_damage_model"]:
            sigma_B = sqrt(
                self.simulation["sensitivity_coefficient"]
                * self.microscope.beam.electrons_per_angstrom
                * (index + 1)
            )
        else:
            sigma_B = 0

        # Either slice or don't
        # if True:  # len(extractor) == 1:

        # Set the atoms in the input after translating them for the offset
        # zslice = extractor[0]
        atoms = self.sample.get_atoms_in_fov(x0, x1)
        logger.info("Simulating with %d atoms" % atoms.data.shape[0])
        # logger.info(
        #     "    Simulating z slice %f -> %f with %d atoms"
        #     % (zslice.x_min[2], zslice.x_max[2], zslice.atoms.data.shape[0])
        # )

        # Set atom sigma
        atoms.data["sigma"] = sigma_B

        if len(atoms.data) > 0:
            coords = atoms.data[["x", "y", "z"]].to_numpy()
            coords = (
                Rotation.from_rotvec((0, angle * pi / 180, 0)).apply(
                    coords - self.sample.centre
                )
                + self.sample.centre
                - (shiftx, shifty + position, 0)
            ).astype("float32")
            atoms.data["x"] = coords[:, 0]
            atoms.data["y"] = coords[:, 1]
            atoms.data["z"] = coords[:, 2]
        # atoms.data = atoms.data.append(parakeet.sample.AtomData(atomic_number=[1,1], x=[750,750], y=[750,750], z=[0,4000],sigma=[0,0],occupancy=[1,1],charge=[0,0]).data)

        input_multislice.spec_atoms = atoms.translate(
            (offset - origin[0], offset - origin[1], 0)
        ).to_multem()
        logger.info("   Got spec atoms")

        if self.simulation["ice"] == True:

            # Create the masker
            masker = multem.Masker(input_multislice.nx, input_multislice.ny, pixel_size)

            # Get the sample centre
            shape = self.sample.shape
            centre = self.sample.centre
            centre = (
                centre[0] + offset - shiftx - origin[0],
                centre[1] + offset - shifty - position - origin[1],
                centre[2],
            )

            # Set the shape
            if shape["type"] == "cube":
                length = shape["cube"]["length"]
                masker.set_cuboid(
                    (
                        centre[0] - length / 2,
                        centre[1] - length / 2,
                        centre[2] - length / 2,
                    ),
                    (length, length, length),
                )
            elif shape["type"] == "cuboid":
                length_x = shape["cuboid"]["length_x"]
                length_y = shape["cuboid"]["length_y"]
                length_z = shape["cuboid"]["length_z"]
                masker.set_cuboid(
                    (
                        centre[0] - length_x / 2,
                        centre[1] - length_y / 2,
                        centre[2] - length_z / 2,
                    ),
                    (length_x, length_y, length_z),
                )
            elif shape["type"] == "cylinder":
                radius = shape["cylinder"]["radius"]
                if not isinstance(radius, Iterable):
                    radius = [radius]
                length = shape["cylinder"]["length"]
                offset_x = shape["cylinder"].get("offset_x", [0] * len(radius))
                offset_z = shape["cylinder"].get("offset_z", [0] * len(radius))
                axis = shape["cylinder"].get("axis", (0, 1, 0))
                masker.set_cylinder(
                    (centre[0], centre[1] - length / 2, centre[2]),
                    axis,
                    length,
                    list(radius),
                    list(offset_x),
                    list(offset_z),
                )

            # Rotate
            origin = centre
            masker.set_rotation(origin, (0, angle * pi / 180.0, 0))

            # Run the simulation
            output_multislice = multem.simulate(system_conf, input_multislice, masker)

        else:

            # Run the simulation
            logger.info("Simulating")
            output_multislice = multem.simulate(system_conf, input_multislice)

        # else:
        #     pass
        # # Slice the specimen atoms
        # def slice_generator(extractor):

        #     # Get the data from the data buffer and return
        #     def prepare(data_buffer):

        #         # Extract the data
        #         atoms = parakeet.sample.AtomData(
        #             data=pandas.concat([d.atoms.data for d in data_buffer])
        #         )
        #         z_min = min([d.x_min[2] for d in data_buffer])
        #         z_max = max([d.x_max[2] for d in data_buffer])
        #         assert z_min < z_max

        #         # Print some info
        #         logger.info(
        #             "    Simulating z slice %f -> %f with %d atoms"
        #             % (z_min, z_max, atoms.data.shape[0])
        #         )

        #         # Cast the atoms
        #         atoms = atoms.translate((offset, offset, 0)).to_multem()

        #         # Return the Z-min, Z-max and atoms
        #         return (z_min, z_max, atoms)

        #     # Loop through the slices and gather atoms until we have more
        #     # than the maximum buffer size. There seems to be an overhead
        #     # to the simulation code so it's better to have as many atoms
        #     # as possible before calling. Doing this is much fast than
        #     # simulating with only a small number of atoms.
        #     max_buffer = 10_000_000
        #     data_buffer = []
        #     for zslice in extractor:
        #         data_buffer.append(zslice)
        #         if sum(d.atoms.data.shape[0] for d in data_buffer) > max_buffer:
        #             yield prepare(data_buffer)
        #             data_buffer = []

        #     # Simulate from the final buffer
        #     if len(data_buffer) > 0:
        #         yield prepare(data_buffer)
        #         data_buffer = []

        # Run the simulation
        # st = time.time()
        # output_multislice = multem.simulate(
        #     system_conf, input_multislice, slice_generator(extractor)
        # )
        # logger.info(
        #     "    Image %d simulated in %d seconds" % (index, time.time() - st)
        # )

        # Get the ideal image data
        # Multem outputs data in column major format. In C++ and Python we
        # generally deal with data in row major format so we must do a
        # transpose here.
        image = numpy.array(output_multislice.data[0].psi_coh).T
        image = image[padding:-padding, padding:-padding]

        # Print some info
        psi_tot = numpy.abs(image) ** 2
        logger.info(
            "Ideal image min/max: %f/%f" % (numpy.min(psi_tot), numpy.max(psi_tot))
        )

        # Compute the image scaled with Poisson noise
        return (index, angle, position, image, (shiftx, shifty))
    def __call__(self, index):
        """
        Simulate a single frame

        Args:
            simulation (object): The simulation object
            index (int): The frame number

        Returns:
            tuple: (angle, image)

        """

        # Get the rotation angle
        angle = self.scan.angles[index]
        position = self.scan.positions[index]

        # The field of view
        nx = self.microscope.detector.nx
        ny = self.microscope.detector.ny
        pixel_size = self.microscope.detector.pixel_size
        margin = self.simulation["margin"]
        x_fov = nx * pixel_size
        y_fov = ny * pixel_size
        offset = margin * pixel_size

        # Get the specimen atoms
        logger.info(f"Simulating image {index+1}")

        # Set the rotation angle
        # input_multislice.spec_rot_theta = angle
        # input_multislice.spec_rot_u0 = simulation.scan.axis

        # x0 = (-offset, -offset)
        # x1 = (x_fov + offset, y_fov + offset)

        # Create the multem system configuration
        system_conf = create_system_configuration(self.device)

        # Create the multem input multislice object
        input_multislice = create_input_multislice(
            self.microscope,
            self.simulation["slice_thickness"],
            self.simulation["margin"],
            "EWRS",
        )

        # Set the specimen size
        input_multislice.spec_lx = x_fov + offset * 2
        input_multislice.spec_ly = y_fov + offset * 2
        input_multislice.spec_lz = numpy.max(self.atoms.data["z"])

        # Set the atoms in the input after translating them for the offset
        input_multislice.spec_atoms = self.atoms.translate(
            (offset, offset, 0)
        ).to_multem()

        # Run the simulation
        output_multislice = multem.simulate(system_conf, input_multislice)

        # Get the ideal image data
        # Multem outputs data in column major format. In C++ and Python we
        # generally deal with data in row major format so we must do a
        # transpose here.
        image = numpy.array(output_multislice.data[0].psi_coh).T

        # Print some info
        psi_tot = numpy.abs(image) ** 2
        logger.info(
            "Ideal image min/max: %f/%f" % (numpy.min(psi_tot), numpy.max(psi_tot))
        )

        # Compute the image scaled with Poisson noise
        return (index, angle, position, image, None)