Skip to content

Probe Engines

PtyLab.ProbeEngines

OPRP

OPRP_storage

Source code in PtyLab/ProbeEngines/OPRP.py 
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
class OPRP_storage:
    def __init__(self, N_probes=5, correct_position=True):
        self.logger = logging.getLogger("OPRP")
        self.N_probes = N_probes
        self.correct_position = correct_position

    def clear(self):
        """Use this when the probe is changed, for instance number of probe modes is changed"""
        for name in ["probes", "probe_indices", "A", "s", "At"]:
            if hasattr(self, name):
                delattr(self, name)

    def push(self, probe, index, N_positions):
        self.xp = getArrayModule(probe)

        probe = probe * self.xp.exp(-1j * self.xp.angle(probe.sum()))

        if not hasattr(self, "probes"):
            self._prepare_probes(probe, N_positions)
            # if it's the first one, make it small so the updates are relatively important
            probe_norm = probe.reshape(self.new_probe_shape)
            probe_norm = probe_norm / self.xp.linalg.norm(probe_norm)
            self.push(probe_norm.reshape(self.original_probe_shape), index, N_positions)
            return

        if self.correct_position:
            probe = self.center_probe(probe, index)

        self.probes[index] = probe.reshape(self.new_probe_shape)
        self.probe_indices[index] = True

    def tsvd(self):
        self.logger.info("Running TSVD")
        if not np.all(self.probe_indices == True):
            indices = self.xp.argwhere(self.probe_indices)
            probes = self.probes[indices]
        else:
            probes = self.probes
        average_power = self.xp.mean(self.xp.abs(probes**2))
        probe_power = self.xp.mean(self.xp.abs(probes**2), axis=-1, keepdims=True)
        probes *= self.xp.mean(average_power / (probe_power + 1e-6))

        A, s, At = self.xp.linalg.svd(probes, full_matrices=False)
        N = self.N_probes
        # calculate effective rank
        pk = s / self.xp.linalg.norm(s.flatten(), ord=1)
        H = -self.xp.sum(pk * self.xp.log(pk))
        eRank = self.xp.exp(H)

        self.A = A[:, :N]
        self.s = s[:N]
        self.At = At[:N]

        self.logger.info(
            f"Effective rank: {eRank}, truncating to {self.N_probes} modes"
        )

        self.logger.info(f"Average displacement: {np.mean(abs(self.center_mass))}")

    def get(self, index):
        """Get the TSVD estimate of the i-th index.

        If the particular index has not been given yet,
        or tsvd has not been run yet, return the averaged probe that we measured so far."""
        if not hasattr(self, "A"):
            # tsvd has not been run yet, return the averaged probe
            return (
                self.probes[self.probe_indices]
                .mean(axis=0)
                .reshape(self.original_probe_shape)
            )

        if self.probe_indices[index]:  # we measured this one, all easy
            A = self.A[index]
        else:  # tsvd has been run, but this particular probe was not given yet
            # beginning, we ask for a probe that we haven't measured yet.
            # In this case, return the typical probe to have some idea
            # Aka set self.A to [1, 0, 0,... 0]
            # we didn't measure it, return the first mode multiplied with the average power
            A = self.A[0].copy()
            A[1:] = 0
            A[0] = 1.0 * self.xp.sign(self.A[0, 0])
        probe = np.matmul(A, self.s[..., None] * self.At)
        # probe = (A * self.s) @ self.At
        # print(probe.shape)
        # move it back
        probe = probe.reshape(self.original_probe_shape)
        if self.correct_position:
            probe = self.uncenter_probe(probe, index)

        return probe

    def center_probe(self, probe, index):
        dpos = (
            np.array(ndimage.center_of_mass(abs(probe**2)))
            - np.array(probe.shape) / 2
        )
        dpos = np.clip(dpos, -2.5, 2.5)
        self.center_mass[index] += 0.01 * dpos
        # move it
        for dim, shift in enumerate(self.center_mass[index]):
            if self.original_probe_shape[dim] != 1:
                shift = np.round(shift).astype(int)
                probe = self.xp.roll(probe, shift=-shift, axis=dim)
        return probe

    def uncenter_probe(self, probe, index):
        probe = probe.reshape(self.original_probe_shape)
        for dim, shift in enumerate(self.center_mass[index]):
            if self.original_probe_shape[dim] != 1:
                shift = np.round(shift).astype(int)
                probe = self.xp.roll(probe, shift=shift, axis=dim)
        return probe

    def _prepare_probes(self, single_probe, N_positions):
        self.original_probe_shape = single_probe.shape
        probe_shape = np.array(single_probe.shape)
        # probe_shape[-2] = single_probe.shape[-1] * single_probe.shape[-2]
        # probe_shape = probe_shape[:-1]
        self.new_probe_shape = np.array([np.product(single_probe.shape)])
        # self.new_probe_shape = probe_shape
        self.N_positions = N_positions
        self.probes = self.xp.zeros(
            (self.N_positions, *self.new_probe_shape), dtype=np.complex64
        )
        self.probe_indices = self.xp.zeros(self.N_positions, dtype=np.bool)

        if self.correct_position:
            shape = (self.N_positions, len(self.original_probe_shape))
            self.center_mass = np.zeros(shape)

    def estimate_CM(self):
        from scipy import ndimage

        for i in range(self.N_positions):
            probe = self.get(i)
            cmass = np.array(ndimage.center_of_mass(abs(probe) ** 2))

            print(i, cmass - np.array(probe.shape) / 2 + 1)
clear()

Use this when the probe is changed, for instance number of probe modes is changed

Source code in PtyLab/ProbeEngines/OPRP.py 
19
20
21
22
23
def clear(self):
    """Use this when the probe is changed, for instance number of probe modes is changed"""
    for name in ["probes", "probe_indices", "A", "s", "At"]:
        if hasattr(self, name):
            delattr(self, name)
get(index)

Get the TSVD estimate of the i-th index.

If the particular index has not been given yet, or tsvd has not been run yet, return the averaged probe that we measured so far.

Source code in PtyLab/ProbeEngines/OPRP.py 
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def get(self, index):
    """Get the TSVD estimate of the i-th index.

    If the particular index has not been given yet,
    or tsvd has not been run yet, return the averaged probe that we measured so far."""
    if not hasattr(self, "A"):
        # tsvd has not been run yet, return the averaged probe
        return (
            self.probes[self.probe_indices]
            .mean(axis=0)
            .reshape(self.original_probe_shape)
        )

    if self.probe_indices[index]:  # we measured this one, all easy
        A = self.A[index]
    else:  # tsvd has been run, but this particular probe was not given yet
        # beginning, we ask for a probe that we haven't measured yet.
        # In this case, return the typical probe to have some idea
        # Aka set self.A to [1, 0, 0,... 0]
        # we didn't measure it, return the first mode multiplied with the average power
        A = self.A[0].copy()
        A[1:] = 0
        A[0] = 1.0 * self.xp.sign(self.A[0, 0])
    probe = np.matmul(A, self.s[..., None] * self.At)
    # probe = (A * self.s) @ self.At
    # print(probe.shape)
    # move it back
    probe = probe.reshape(self.original_probe_shape)
    if self.correct_position:
        probe = self.uncenter_probe(probe, index)

    return probe

StandardProbe

LinearProbe

Source code in PtyLab/ProbeEngines/StandardProbe.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
class LinearProbe:
    def __init__(self):
        self.logger = logging.getLogger("SHG")
        self.probe = None
        self.probe_temp = None

    def clear(self):
        pass

    def push(self, new_probe, index, N_positions, factor=1.0, force=False):
        """
        Set the current estimate of the probe to new_probe.
        """
        if force:
            self.probe = new_probe
        elif self.probe is not None:
            self.probe = new_probe * factor + (1 - factor) * self.probe
        else:
            self.probe = new_probe
        self.probe_temp = self.probe.copy()

    def set_temporary(self, probe):
        """These map to self.reconstruction.probe. Can be used for quick updates in the calculation of the probe.

        Once you're done, make it official by updating with push()"""
        self.probe_temp = probe

    def get_temporary(self):
        return self.probe_temp

    def get(self, index):
        return self.probe

    def roll(self, dy, dx):
        self.probe = self.probe_temp.copy()
        xp = getArrayModule(self.probe)
        self.probe = xp.roll(self.probe, (-dy, -dx), axis=(-2, -1))
        self.probe_temp = self.probe.copy()
push(new_probe, index, N_positions, factor=1.0, force=False)

Set the current estimate of the probe to new_probe.

Source code in PtyLab/ProbeEngines/StandardProbe.py 
16
17
18
19
20
21
22
23
24
25
26
def push(self, new_probe, index, N_positions, factor=1.0, force=False):
    """
    Set the current estimate of the probe to new_probe.
    """
    if force:
        self.probe = new_probe
    elif self.probe is not None:
        self.probe = new_probe * factor + (1 - factor) * self.probe
    else:
        self.probe = new_probe
    self.probe_temp = self.probe.copy()
set_temporary(probe)

These map to self.reconstruction.probe. Can be used for quick updates in the calculation of the probe.

Once you're done, make it official by updating with push()

Source code in PtyLab/ProbeEngines/StandardProbe.py 
28
29
30
31
32
def set_temporary(self, probe):
    """These map to self.reconstruction.probe. Can be used for quick updates in the calculation of the probe.

    Once you're done, make it official by updating with push()"""
    self.probe_temp = probe

SHGProbe

Bases: LinearProbe

Source code in PtyLab/ProbeEngines/StandardProbe.py 
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
class SHGProbe(LinearProbe):
    def __init__(self):
        super().__init__()
        self.logger = logging.getLogger("SHG")
        self.probe = None  # wavelength = wavelength  * self.nonlinearity
        self.nonlinearity = 2

    def clear(self):
        pass

    def push(self, new_probe_nonlinear, index, N_positions):
        """Gets the update of the nonlinear part"""
        if self.probe is not None:
            if isGpuArray(new_probe_nonlinear) and not isGpuArray(self.probe):
                import cupy as cp

                print(" Moving self.probe to GPU")
                self.probe = cp.array(self.probe)
            else:
                xp = getArrayModule(self.probe)
                new_probe_nonlinear = xp.array(new_probe_nonlinear)
        else:
            xp = getArrayModule(new_probe_nonlinear)

        # what's actually pushed is the second harmonic. We need to update that

        if self.probe is None:
            self.probe = new_probe_nonlinear * 0
        # try "newtons" method

        # Solve for the new estimate, and update the original estimate based on it
        new_probe_estimate = new_probe_nonlinear ** (1.0 / self.nonlinearity)

        diff = new_probe_estimate - self.probe
        self.probe += diff / (2 * self.nonlinearity)
        # diff = new_probe_nonlinear - self.probe ** self.nonlinearity
        if N_positions == -1:
            print(np.linalg.norm(self.probe**self.nonlinearity - new_probe_nonlinear))
        # update_fundamental = self.nonlinearity * diff
        # self.probe += update_fundamental
        self.probe_temp = self.probe.copy() ** self.nonlinearity

    def change_nonlinearity(self, nonlinearity):
        last_probe = self.get(None).copy()
        self.nonlinearity = nonlinearity
        self._push_hard(last_probe)

    def _push_hard(self, new_probe, number_of_iterations=50):
        xp = getArrayModule(self.probe)
        new_probe = xp.array(new_probe)
        for i in range(number_of_iterations):
            self.push(new_probe, None, -1)
            print(xp.linalg.norm(self.get(None) - new_probe))

    def get(self, index):
        return self.probe**self.nonlinearity

    def get_fundamental(self):
        return self.probe
push(new_probe_nonlinear, index, N_positions)

Gets the update of the nonlinear part

Source code in PtyLab/ProbeEngines/StandardProbe.py 
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
def push(self, new_probe_nonlinear, index, N_positions):
    """Gets the update of the nonlinear part"""
    if self.probe is not None:
        if isGpuArray(new_probe_nonlinear) and not isGpuArray(self.probe):
            import cupy as cp

            print(" Moving self.probe to GPU")
            self.probe = cp.array(self.probe)
        else:
            xp = getArrayModule(self.probe)
            new_probe_nonlinear = xp.array(new_probe_nonlinear)
    else:
        xp = getArrayModule(new_probe_nonlinear)

    # what's actually pushed is the second harmonic. We need to update that

    if self.probe is None:
        self.probe = new_probe_nonlinear * 0
    # try "newtons" method

    # Solve for the new estimate, and update the original estimate based on it
    new_probe_estimate = new_probe_nonlinear ** (1.0 / self.nonlinearity)

    diff = new_probe_estimate - self.probe
    self.probe += diff / (2 * self.nonlinearity)
    # diff = new_probe_nonlinear - self.probe ** self.nonlinearity
    if N_positions == -1:
        print(np.linalg.norm(self.probe**self.nonlinearity - new_probe_nonlinear))
    # update_fundamental = self.nonlinearity * diff
    # self.probe += update_fundamental
    self.probe_temp = self.probe.copy() ** self.nonlinearity