Skip to content

Commit

Permalink
with the instruction on how to recover the corrected data
Browse files Browse the repository at this point in the history
  • Loading branch information
shcrela committed May 21, 2023
1 parent 72e08dc commit c48f7ee
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 21 deletions.
42 changes: 21 additions & 21 deletions UntiltPlane.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,53 +7,53 @@

class UntiltPlane(object):
"""Level the plane if the sample is tilted.
Left click on the image to select the points which will be used to create
the offset plane. On right-click, this plane will be substracted from
your image's original data, and the image will be redrawn with updated values.
Attention: The resulting data will always have zero as its' minimal value.
That is because after substracting the plane, we also substract the minimum
of that result.
Important Note: There is a strong assumption that the offset plane is linear.
Parameters:
-----------
data: 2D numpy.ndarray
Your input data. NaNs are fine.
Your input data. NaNs are fine.
a, b, c: floats
The initial parameters of the flat offset plane.
Defaults: `a = np.nanmin(data)`, `b = 0`, `c = 0`
ftol: float
What precision of the fit is satisfactory? Default is 1e-10
**kwargs:
Arguments passed on to matplotlib.figure
Output:
----------
corrected_data: numpy.ndarray
The updated ("untilted") values of your input array.
plane_params: tuple(float, float, float)
The tuple of fitted plane parameters (a, b, c). You can also recover them
as `a`, `b` and `c`
Example:
---------
```
# You just need to run something like:
>>>my_untilter = UntiltPlane(data)
# Then, if you want to recover the corrected data:
>>>corrected_data = my.untilter.corrected_data
>>>corrected_data = my_untilter.corrected_data
# If you wish to reconstruct the offset plane:
>>>a, b, c = my_untilter.plan_params
>>>offset_plane = np.fromfunction(lambda y, x: a + b*x + c*y, data.shape)
```
"""

def __init__(self, data, a="min", b=0, c=0, ftol=1e-10, **kwargs):
self.data = data

self.data = data
if a == "min":
self.a = np.nanmin(self.data)
else:
Expand All @@ -66,12 +66,12 @@ def __init__(self, data, a="min", b=0, c=0, ftol=1e-10, **kwargs):
self.plane_data = []
self.my_points = []


figsize = kwargs.pop("figsize", (12, 8))
facecolor = kwargs.pop("facecolor", "oldlace")
self.fig, self.aximg = plt.subplots(figsize=figsize, facecolor=facecolor,
**kwargs)

self.img = self.aximg.imshow(self.data, cmap="cividis")
self.aximg.tick_params(top=True, labeltop=True, bottom=False, labelbottom=False)
self.aximg.set_title("Left-click to assign the points that will define\n"+
Expand All @@ -82,7 +82,7 @@ def __init__(self, data, a="min", b=0, c=0, ftol=1e-10, **kwargs):

def flat_plane_z(self, y, x):
return self.a + self.b*x + self.c*y

def optimize_flat_plane_z(self, independent_coords, a, b, c):
x = independent_coords[:, 0]
y = independent_coords[:, 1]
Expand All @@ -105,15 +105,15 @@ def onclick(self, event):
elif event.button != 1:
self.fig.canvas.mpl_disconnect(self.cid)
self.detilt()

def detilt(self):
"""Remove the points, fit the (linear) plane, and update the image"""

# Remove all the points:
for point in self.my_points:
point.remove()
self.my_points = []

clicked_data = np.array(self.plane_data)
independent_coords = clicked_data[:, :2]
measured_z = clicked_data[:, 2]
Expand All @@ -127,18 +127,18 @@ def detilt(self):
initial_guess,
method="lm",
ftol=self.ftol)


self.a, self.b, self.c = self.plane_params
base_plane = np.fromfunction(self.flat_plane_z, self.data.shape)
self.corrected_data = self.data - base_plane
# Correct the offset (not sure whether we should do this or not)
self.corrected_data -= np.nanmin(self.corrected_data)

im_min, im_max = 0, np.nanmax(self.corrected_data)
gs = gridspec.GridSpec(1,2)
self.aximg.set_position(gs[0].get_position(self.fig))
self.aximg.set_subplotspec(gs[0])
self.aximg.set_subplotspec(gs[0])
self.axcorr = self.fig.add_subplot(gs[1])
self.axcorr.imshow(self.corrected_data, cmap="cividis")
self.axcorr.tick_params(top=True, labeltop=True, bottom=False, labelbottom=False)
Expand Down
5 changes: 5 additions & 0 deletions test.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,3 +16,8 @@
raw = np.flipud(np.loadtxt(filename, delimiter=','))

my_plane = UntiltPlane(raw)

# %%
######## Run the line below only afer you've finished the above ########
######## (shift + enter to run cell by cell) ########
my_corrected_data = my_plane.corrected_data

0 comments on commit c48f7ee

Please sign in to comment.