def plot_spectrum_to_ax(ax: Axes, spectrum: Spectrum1D, title: str,
c_range: SpectralRegion = None, h_range: SpectralRegion = None,
sky_flux: [Quantity] = None, nss_spec_flux: [Quantity] = None):
ax.set_xlabel(f"Wavelength [{spectrum.wavelength.unit}]")
ax.set_ylabel(f"flux [adu]")
ax.set_title(title)
ax.set_xticks(_calculate_ticks(spectrum.wavelength.value, 250), minor=False)
ax.set_xticks(_calculate_ticks(spectrum.wavelength.value, 50), minor=True)
ax.set_xticklabels(ax.get_xticks(minor=False))
color = "b" if min(ax.get_xticks(minor=True)) < 4500 else "r"
wavelength = spectrum.wavelength
ax.plot(wavelength, spectrum.flux, color=color, linestyle="-", linewidth=0.25)
if nss_spec_flux is not None and len(nss_spec_flux) == len(wavelength):
ax.plot(wavelength, nss_spec_flux, color="grey", linestyle="-", linewidth=0.25, alpha=0.3)
if sky_flux is not None and len(sky_flux) == len(wavelength):
ax.plot(wavelength, sky_flux, color="c", linestyle="-", linewidth=0.25, alpha=0.5)
if c_range is not None:
ax.axvspan(xmin=c_range.lower.value, xmax=c_range.upper.value, color="c", alpha=0.05)
if h_range is not None:
ax.axvspan(xmin=h_range.lower.value, xmax=h_range.upper.value, color=color, alpha=0.05)
return
def highlight_weekends(dates: [datetime], ax: Axes):
# first
if dates[0].weekday() == 5 or dates[0].weekday() == 6:
start = dates[0]
else:
start = dates[0] + timedelta(5 - dates[0].weekday())
start = start.replace(hour=0, minute=0, second=0, tzinfo=pytz.utc)
monday = start + timedelta(7 - start.weekday())
monday = monday.replace(hour=0, minute=0, second=0, tzinfo=pytz.utc)
ax.axvspan(start, monday, facecolor='green', edgecolor='none', alpha=0.1)
# middle and last
while True:
saturday = monday + timedelta(5)
if saturday > dates[-1]: # no more weekends
break
monday = saturday + timedelta(2)
if monday > dates[-1]: # log ends during weekend
ax.axvspan(saturday, dates[-1], facecolor='green', edgecolor='none', alpha=0.1)
break
else:
ax.axvspan(saturday, monday, facecolor='green', edgecolor='none', alpha=0.1)
def plot_arrows(
self,
symbol: str,
ax: Axes,
start: datetime = None,
end: datetime = None,
) -> None:
pa: ProfitAttribution
pas = sorted(
[pa for pa in self.pas if pa.sym == symbol],
key=lambda pa: pa.open_trade.time,
)
ax.plot(
[date2num(pa.buy_trade.time) for pa in pas],
[pa.buy_trade.price for pa in pas],
color="#0008",
marker="2",
lw=0,
markersize=15,
zorder=-1,
)
ax.plot(
[date2num(pa.sell_trade.time) for pa in pas],
[pa.sell_trade.price for pa in pas],
color="#0008",
marker="1",
lw=0,
markersize=15,
zorder=-1,
)
if not pas:
return
max_qty = max([pa.qty for pa in pas])
for pa in pas:
x0 = date2num(pa.open_trade.time)
x1 = date2num(pa.close_trade.time)
y0 = pa.open_trade.price
y1 = pa.close_trade.price
if pa.net_gain > 0:
color = "green"
else:
color = "red"
ax.plot(
[x0, x1],
[y0, y1],
lw=(w := 0.5 + 2.5 * pa.qty / max_qty),
zorder=1000 - w,
color=color,
)
ax.figure.autofmt_xdate()
if start is None:
lbx = min([pa.start_time for pa in pas]) - ONE_DAY
else:
lbx = start
if end is None:
ubx = max([pa.end_time for pa in pas]) + ONE_DAY
else:
ubx = end
ax.set_xlim(lbx, ubx)
lby = max(min([min(pa.buy_price, pa.sell_price) for pa in pas]) - 1, 0)
uby = max([max(pa.buy_price, pa.sell_price) for pa in pas]) + 1
ax.set_ylim(lby, uby)
ax.set_xlabel("Trade dt")
ax.set_ylabel("Trade price")
for day in date_range(
lbx.replace(hour=0, minute=0, second=0), ubx, freq="1D"
):
if day.weekday() in (5, 6):
continue
ax.axvspan(
day.replace(hour=9, minute=30),
day.replace(hour=16),
0,
1.0,
facecolor="#0000ff20",
zorder=-1e3,
)
buys = np.array(
[pa.buy_price for pa in pas for _ in range(abs(pa.qty))]
)
sells = np.array(
[pa.sell_price for pa in pas for _ in range(abs(pa.qty))]
)
total_gain = (sells - np.minimum(sells, buys)).sum()
total_loss = (sells - np.maximum(sells, buys)).sum()
ax.set_title(
f"{symbol} profits: {total_gain:.0f} - {-total_loss:.0f} = "
f"{total_gain + total_loss:.0f}"
)
ax.xaxis.set_major_locator(WeekdayLocator(byweekday=MO, tz=TZ_EASTERN))
ax.xaxis.set_minor_locator(DayLocator(tz=TZ_EASTERN))
ax.xaxis.set_minor_formatter(DateFormatter("%d"))
ax.grid(color="#808080", lw=0.5)
ax.grid(color="#808080", lw=0.25, axis="x", which="minor")
