diff --git a/src/shadowfinder/shadowfinder.py b/src/shadowfinder/shadowfinder.py index 63f67ce..ad7637b 100644 --- a/src/shadowfinder/shadowfinder.py +++ b/src/shadowfinder/shadowfinder.py @@ -1,9 +1,11 @@ from pytz import timezone, utc +import datetime import pandas as pd from suncalc import get_position import numpy as np import matplotlib.pyplot as plt import matplotlib.colors as colors +from matplotlib.figure import Figure from mpl_toolkits.basemap import Basemap from timezonefinder import TimezoneFinder import json @@ -20,6 +22,10 @@ def __init__( time_format="utc", sun_altitude_angle=None, ): + self.sun_altitude_angle = None + self.object_height = None + self.shadow_length = None + self.set_details( date_time, object_height, shadow_length, time_format, sun_altitude_angle ) @@ -39,15 +45,15 @@ def __init__( self.min_lon = -180 self.max_lon = 180 - def set_details( - self, - date_time, - object_height=None, - shadow_length=None, - time_format=None, - sun_altitude_angle=None, - ): - + def set_datetime( + self, date_time: datetime.datetime = None, time_format=None + ) -> None: + if time_format is not None: + assert time_format in [ + "utc", + "local", + ], "time_format must be 'utc' or 'local'" + self.time_format = time_format if date_time is not None and date_time.tzinfo is not None: warn( "date_time is expected to be timezone naive (i.e. tzinfo=None). Any timezone information will be ignored." @@ -55,12 +61,15 @@ def set_details( date_time = date_time.replace(tzinfo=None) self.date_time = date_time - if time_format is not None: - assert time_format in [ - "utc", - "local", - ], "time_format must be 'utc' or 'local'" - self.time_format = time_format + def set_details( + self, + date_time, + object_height=None, + shadow_length=None, + time_format=None, + sun_altitude_angle=None, + ) -> None: + self.set_datetime(date_time=date_time, time_format=time_format) # height and length must have the same None-ness # either height or angle must be set (but not both or neither) @@ -93,14 +102,14 @@ def set_details( # Lengths and angle are None and we use the same values as before pass - def quick_find(self, timezone_grid="timezone_grid.json"): + def quick_find(self, timezone_grid="timezone_grid.json") -> None: # try to load timezone grid from file, generate if not found try: self.load_timezone_grid(timezone_grid) except FileNotFoundError: self.generate_timezone_grid() - self.find_shadows() + _ = self.find_shadows() fig = self.plot_shadows() if self.sun_altitude_angle is not None: @@ -110,7 +119,7 @@ def quick_find(self, timezone_grid="timezone_grid.json"): fig.savefig(file_name) - def generate_timezone_grid(self): + def generate_timezone_grid(self) -> None: lats = np.arange(self.min_lat, self.max_lat, self.angular_resolution) lons = np.arange(self.min_lon, self.max_lon, self.angular_resolution) @@ -124,7 +133,7 @@ def generate_timezone_grid(self): ] ) - def save_timezone_grid(self, filename="timezone_grid.json"): + def save_timezone_grid(self, filename="timezone_grid.json") -> None: data = { "min_lat": self.min_lat, "max_lat": self.max_lat, @@ -136,7 +145,7 @@ def save_timezone_grid(self, filename="timezone_grid.json"): json.dump(data, open(filename, "w")) - def load_timezone_grid(self, filename="timezone_grid.json"): + def load_timezone_grid(self, filename="timezone_grid.json") -> None: data = json.load(open(filename, "r")) self.min_lat = data["min_lat"] @@ -151,14 +160,11 @@ def load_timezone_grid(self, filename="timezone_grid.json"): self.lons, self.lats = np.meshgrid(lons, lats) self.timezones = np.array(data["timezones"]) - def find_shadows(self): - # Evaluate the sun's length at a grid of points on the Earth's surface - - if self.lats is None or self.lons is None or self.timezones is None: - self.generate_timezone_grid() + def generate_valid_datetimes_lats_lons_and_mask(self, timestamp): + mask = None if self.time_format == "utc": - valid_datetimes = utc.localize(self.date_time) + valid_datetimes = utc.localize(timestamp) valid_lats = self.lats.flatten() valid_lons = self.lons.flatten() elif self.time_format == "local": @@ -168,7 +174,7 @@ def find_shadows(self): None if tz is None else timezone(tz) - .localize(self.date_time) + .localize(timestamp) .astimezone(utc) .timestamp() ) @@ -187,6 +193,29 @@ def find_shadows(self): # Convert the datetimes to pandas series of timestamps valid_datetimes = pd.to_datetime(valid_datetimes, unit="s", utc=True) + return valid_datetimes, valid_lons, valid_lats, mask + + def find_shadows( + self, + object_height: float = None, + shadow_length: float = None, + sun_altitude_angle: float = None, + timestamp: datetime.datetime = None, + ) -> np.ndarray: + # Evaluate the sun's length at a grid of points on the Earth's surface + + object_height = object_height if object_height else self.object_height + shadow_length = shadow_length if shadow_length else self.shadow_length + sun_altitude_angle = ( + sun_altitude_angle if sun_altitude_angle else self.sun_altitude_angle + ) + timestamp = timestamp if timestamp else self.date_time + if self.lats is None or self.lons is None or self.timezones is None: + self.generate_timezone_grid() + + valid_datetimes, valid_lons, valid_lats, mask = ( + self.generate_valid_datetimes_lats_lons_and_mask(timestamp=timestamp) + ) pos_obj = get_position(valid_datetimes, valid_lons, valid_lats) valid_sun_altitudes = pos_obj["altitude"] # in radians @@ -194,9 +223,9 @@ def find_shadows(self): # If object height and shadow length are set the sun altitudes are used # to calculate the shadow lengths across the world and then compared to # the expected shadow length. - if self.object_height is not None and self.shadow_length is not None: + if object_height is not None and shadow_length is not None: # Calculate the shadow length - shadow_lengths = self.object_height / np.apply_along_axis( + shadow_lengths = object_height / np.apply_along_axis( np.tan, 0, valid_sun_altitudes ) @@ -204,17 +233,15 @@ def find_shadows(self): shadow_lengths[valid_sun_altitudes <= 0] = np.nan # Show the relative difference between the calculated shadow length and the observed shadow length - location_likelihoods = ( - shadow_lengths - self.shadow_length - ) / self.shadow_length + location_likelihoods = (shadow_lengths - shadow_length) / shadow_length # If the sun altitude angle is set then this value is directly compared # to the sun altitudes across the world. - elif self.sun_altitude_angle is not None: + elif sun_altitude_angle is not None: # Show relative difference between sun altitudes location_likelihoods = ( - np.array(valid_sun_altitudes) - radians(self.sun_altitude_angle) - ) / radians(self.sun_altitude_angle) + np.array(valid_sun_altitudes) - radians(sun_altitude_angle) + ) / radians(sun_altitude_angle) # Replace points where the sun is below the horizon location_likelihoods[valid_sun_altitudes <= 0] = np.nan @@ -237,12 +264,41 @@ def find_shadows(self): self.location_likelihoods = np.reshape( self.location_likelihoods, np.shape(self.lons), order="A" ) + return self.location_likelihoods + + def get_plot_data(self, data: dict): + object_height = data.get("object_height", self.object_height) + shadow_length = data.get("shadow_length", self.shadow_length) + sun_altitude_angle = data.get("sun_altitude_angle", self.sun_altitude_angle) + timestamp = data.get("timestamp", self.date_time) + location_likelihoods = data.get( + "location_likelihoods", self.location_likelihoods + ) + plt_title = data.get("plt_title", None) + return ( + object_height, + shadow_length, + sun_altitude_angle, + timestamp, + location_likelihoods, + plt_title, + ) def plot_shadows( self, figure_args={"figsize": (12, 6)}, basemap_args={"projection": "cyl", "resolution": "c"}, - ): + plot_data: dict = None, + ) -> Figure: + plot_data = plot_data if plot_data else {} + ( + object_height, + shadow_length, + sun_altitude_angle, + timestamp, + location_likelihoods, + plt_title, + ) = self.get_plot_data(data=plot_data) fig = plt.figure(**figure_args) @@ -262,19 +318,75 @@ def plot_shadows( m.pcolormesh( x, y, - np.abs(self.location_likelihoods), + np.abs(location_likelihoods), cmap=cmap, norm=norm, alpha=0.7, ) # plt.colorbar(label='Relative Shadow Length Difference') - - if self.sun_altitude_angle is not None: - plt_title = f"Possible Locations at {self.date_time.strftime('%Y-%m-%d %H:%M:%S')} {self.time_format.title()}\n(sun altitude angle: {self.sun_altitude_angle})" - else: - plt_title = f"Possible Locations at {self.date_time.strftime('%Y-%m-%d %H:%M:%S')} {self.time_format.title()}\n(object height: {self.object_height}, shadow length: {self.shadow_length})" + if plt_title is None: + if sun_altitude_angle is not None: + plt_title = f"Possible Locations at {timestamp.strftime('%Y-%m-%d %H:%M:%S')} {self.time_format.title()}\n(sun altitude angle: {sun_altitude_angle})" + else: + plt_title = f"Possible Locations at {timestamp.strftime('%Y-%m-%d %H:%M:%S')} {self.time_format.title()}\n(object height: {object_height}, shadow length: {shadow_length})" plt.title(plt_title) self.fig = fig return fig + + def find_multiple_shadows( + self, + object_heights: list[float], + shadow_lengths: list[float], + timestamps: list[datetime.datetime], + ) -> Figure: + """ + Find possible locations using multiple shadows as input. + + If the same object is used every time, the object height can be provided as a list with a single value + """ + if len(shadow_lengths) != len(timestamps): + raise ValueError( + f"Argument lists are expected to have the same length. Provided {len(shadow_lengths)},{len(timestamps)}." + ) + elif len(object_heights) != len(timestamps): + if len(object_heights) == 1: + object_heights = [object_heights[0] for _ in timestamps] + else: + raise ValueError( + f"object_heights must be a list of the same length as the shadow lengths and timestamps, or contain only 1 value" + ) + if self.lats is None or self.lons is None or self.timezones is None: + self.generate_timezone_grid() + + arr_location_likelihoods = np.zeros( + (len(timestamps), self.lons.shape[0], self.lons.shape[1]) + ) + for k, (obj_height, shadow_length, timestamp) in enumerate( + zip(object_heights, shadow_lengths, timestamps) + ): + arr_location_likelihoods[k, :, :] = self.find_shadows( + object_height=obj_height, + shadow_length=shadow_length, + timestamp=timestamp, + ) + + # merge/logic for combining the shadows. #TODO + location_likelihoods = np.max(arr_location_likelihoods, axis=0) + + plt_title = f"Possible Locations for measurements at [" + for ts in timestamps: + plt_title += f"{ts.strftime('%Y-%m-%d %H:%M:%S')}, " + plt_title = plt_title[:-2] + f"] ({self.time_format.title()}\n)" + + plot_data = { + "location_likelihoods": location_likelihoods, + "plt_title": plt_title, + "timestamps": timestamps, + "object_heights": object_heights, + "shadow_lengths": shadow_lengths, + } + fig = self.plot_shadows(plot_data=plot_data) + + return fig diff --git a/tests/test_shadowfinder.py b/tests/test_shadowfinder.py index 531c61e..3460d28 100644 --- a/tests/test_shadowfinder.py +++ b/tests/test_shadowfinder.py @@ -1,4 +1,5 @@ -from datetime import datetime +from datetime import datetime, timedelta + from shadowfinder import ShadowFinder @@ -15,3 +16,15 @@ def test_creation_with_valid_arguments_should_pass(): ShadowFinder( object_height=object_height, shadow_length=shadow_length, date_time=date_time ) + +def test_find_multiple_shadows(): + """Test find_multiple_shadows""" + # GIVEN + object_heights = [6,6,6,6] + shadow_lengths = [3.2,3.0,2.9,2.8] + time_offsets = range(4) + timestamps = [datetime.now()+timedelta(hours=f) for f in time_offsets] + + shadow_finder = ShadowFinder() + # WHEN / THEN + figure = shadow_finder.find_multiple_shadows(object_heights=object_heights, shadow_lengths=shadow_lengths, timestamps=timestamps)