Auto-update: Fri Nov 15 13:05:58 PST 2024

This commit is contained in:
sanj 2024-11-15 13:05:58 -08:00
parent 2d49d296fb
commit f44d2e7122

View file

@ -127,57 +127,89 @@ async def generate_and_save_heatmap(
end_date: Optional[Union[str, int, datetime]] = None, end_date: Optional[Union[str, int, datetime]] = None,
output_path: Optional[Path] = None output_path: Optional[Path] = None
) -> Path: ) -> Path:
""" """
Generate a heatmap for the given date range and save it as a PNG file using matplotlib. Generate a heatmap for the given date range and save it as a PNG file.
:param start_date: The start date for the map (or the only date if end_date is not provided) :param start_date: The start date for the map (or the only date if end_date is not provided)
:param end_date: The end date for the map (optional) :param end_date: The end date for the map (optional)
:param output_path: The path to save the PNG file (optional) :param output_path: The path to save the PNG file (optional)
:return: The path where the PNG file was saved :return: The path where the PNG file was saved
""" """
try: try:
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
import numpy as np import contextily as ctx
import numpy as np
start_date = await dt(start_date) from matplotlib.colors import LinearSegmentedColormap
if end_date:
end_date = await dt(end_date) start_date = await dt(start_date)
else: if end_date:
end_date = start_date.replace(hour=23, minute=59, second=59) end_date = await dt(end_date)
else:
locations = await fetch_locations(start_date, end_date) end_date = start_date.replace(hour=23, minute=59, second=59)
if not locations:
raise ValueError("No locations found for the given date range") locations = await fetch_locations(start_date, end_date)
if not locations:
lats = [loc.latitude for loc in locations] raise ValueError("No locations found for the given date range")
lons = [loc.longitude for loc in locations]
lats = np.array([loc.latitude for loc in locations])
plt.style.use('dark_background') lons = np.array([loc.longitude for loc in locations])
fig, ax = plt.subplots(figsize=(10, 6))
# Calculate bounds with 5% buffer
# Create heatmap lat_range = max(lats) - min(lats)
heatmap, xedges, yedges = np.histogram2d(lons, lats, bins=50) lon_range = max(lons) - min(lons)
extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]] buffer = max(lat_range, lon_range) * 0.05
# Plot with no axes or labels # Enforce minimum zoom
ax.imshow(heatmap.T, extent=extent, origin='lower', cmap='hot', interpolation='gaussian') MIN_RANGE = 0.05 # roughly 3-4 miles
ax.axis('off') lat_range = max(lat_range, MIN_RANGE)
lon_range = max(lon_range, MIN_RANGE)
# Remove white border
plt.gca().set_position([0, 0, 1, 1]) bounds = [
min(lons) - buffer,
if output_path is None: max(lons) + buffer,
output_path, relative_path = assemble_journal_path(end_date, filename="map", extension=".png", no_timestamp=True) min(lats) - buffer,
max(lats) + buffer
plt.savefig(output_path, bbox_inches='tight', pad_inches=0, transparent=True) ]
plt.close()
# Create figure with fixed size
l.info(f"Heatmap saved as PNG: {output_path}") fig, ax = plt.subplots(figsize=(6.4, 3.6), dpi=100) # 640x360 pixels
return output_path
# Add dark basemap
except Exception as e: ctx.add_basemap(
l.error(f"Error generating heatmap: {str(e)}") ax,
raise crs='EPSG:4326',
source=ctx.providers.CartoDB.DarkMatter,
zoom='auto',
bbox=bounds
)
# Create heatmap overlay
heatmap = ax.hexbin(
lons, lats,
extent=bounds,
gridsize=25,
cmap='hot',
alpha=0.6,
zorder=2
)
# Remove axes and margins
ax.set_axis_off()
plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
if output_path is None:
output_path, relative_path = assemble_journal_path(end_date, filename="map", extension=".png", no_timestamp=True)
plt.savefig(output_path, bbox_inches='tight', pad_inches=0, dpi=100)
plt.close()
l.info(f"Heatmap saved as PNG: {output_path}")
return output_path
except Exception as e:
l.error(f"Error generating heatmap: {str(e)}")
raise
async def generate_map(start_date: datetime, end_date: datetime, max_points: int): async def generate_map(start_date: datetime, end_date: datetime, max_points: int):