class NumExprSample2(NumExprABC):
inputs = [
SlotDescriptor("first", type=Table, required=True),
SlotDescriptor("second", type=Table, required=True),
]
outputs = [SlotDescriptor("table", type=Table, required=False)]
expr = {
"_1:float64": "{first._2}+2*{second._3}",
"_2:float64": "{first._3}-5*{second._2}",
}
class TestModule(Module):
inputs = [SlotDescriptor("a"), SlotDescriptor("b", required=False)]
outputs = [SlotDescriptor("c"), SlotDescriptor("d", required=False)]
def __init__(self, **kwds: Any) -> None:
super(TestModule, self).__init__(**kwds)
def run_step(self, run_number: int, step_size: int,
howlong: float) -> ReturnRunStep:
return self._return_run_step(self.state_blocked, 0)
class MixUfuncCustomBinary(MixUfuncABC):
inputs = [
SlotDescriptor("first", type=Table, required=True),
SlotDescriptor("second", type=Table, required=True),
]
outputs = [SlotDescriptor("table", type=Table, required=False)]
expr = {
"_1:float64": (custom_binary_ufunc, "first._2", "second._3"),
"_2:float64": (np.log, "second._3"),
}
class MixUfuncSample2(MixUfuncABC):
inputs = [
SlotDescriptor("first", type=Table, required=True),
SlotDescriptor("second", type=Table, required=True),
]
outputs = [SlotDescriptor("result", type=Table, required=False)]
expr = {
"_1:float64": (np.add, "first._2", "second._3"),
"_2:float64": (np.log, "second._3"),
}
class NumExprSample(NumExprABC):
inputs = [
SlotDescriptor("first", type=Table, required=True),
SlotDescriptor("second", type=Table, required=True),
]
outputs = [
SlotDescriptor("result",
type=Table,
required=False,
datashape={"first": ["_1", "_2"]})
]
expr = {"_1": "{first._2}+2*{second._3}", "_2": "{first._3}-5*{second._2}"}
class MixUfuncSample(MixUfuncABC):
inputs = [
SlotDescriptor("first", type=Table, required=True),
SlotDescriptor("second", type=Table, required=True),
]
outputs = [
SlotDescriptor("result",
type=Table,
required=False,
datashape={"first": ["_1", "_2"]})
]
expr = {
"_1": (np.add, "first._2", "second._3"),
"_2": (np.log, "second._3")
}
class Heatmap(TableModule):
"Heatmap module"
parameters = [
("cmax", np.dtype(float), np.nan),
("cmin", np.dtype(float), np.nan),
("high", np.dtype(int), 65536),
("low", np.dtype(int), 0),
("filename", np.dtype(object), None),
("history", np.dtype(int), 3),
]
inputs = [SlotDescriptor("array", type=Table)]
# schema = [('image', np.dtype(object), None),
# ('filename', np.dtype(object), None),
# UPDATE_COLUMN_DESC]
schema = "{filename: string, time: int64}"
def __init__(self, colormap: None = None, **kwds: Any) -> None:
super(Heatmap, self).__init__(output_required=False, **kwds)
self.tags.add(self.TAG_VISUALIZATION)
self.colormap = colormap
self.default_step_size = 1
name = self.generate_table_name("Heatmap")
self.result = Table(name, dshape=Heatmap.schema, create=True)
def predict_step_size(self, duration: float) -> int:
_ = duration
# Module sample is constant time (supposedly)
return 1
def run_step(self, run_number: int, step_size: int,
howlong: float) -> ReturnRunStep:
dfslot = self.get_input_slot("array")
input_df = dfslot.data()
dfslot.deleted.next()
indices = dfslot.created.next()
steps = indices_len(indices)
if steps == 0:
indices = dfslot.updated.next()
steps = indices_len(indices)
if steps == 0:
return self._return_run_step(self.state_blocked, steps_run=1)
histo = input_df.last()["array"]
if histo is None:
return self._return_run_step(self.state_blocked, steps_run=1)
params = self.params
cmax: Optional[float] = cast(float, params.cmax)
assert cmax
if np.isnan(cmax):
cmax = None
cmin: Optional[float] = cast(float, params.cmin)
assert cmin
if np.isnan(cmin):
cmin = None
high: int = cast(int, params.high)
low: int = cast(int, params.low)
try:
if cmin is None:
cmin = histo.min()
if cmax is None:
cmax = histo.max()
# cscale = cmax - cmin
scale_hl = float(high - low)
# scale = float(high - low) / cscale
# data = (sp.special.cbrt(histo) * 1.0 - cmin) * scale + 0.4999
data = (sp.special.cbrt(histo) * 1.0 - cmin) * scale_hl + 0.4999
data[data > high] = high
data[data < 0] = 0
data = np.cast[np.uint32](data)
if low != 0:
data += low
image = Image.fromarray(data, mode="I")
image = image.transpose(Image.FLIP_TOP_BOTTOM)
filename = params.filename
except Exception:
image = None
filename = None
if filename is not None:
try:
if re.search(r"%(0[\d])?d", filename):
filename = filename % (run_number)
filename = self.storage.fullname(self, filename)
# TODO should do it atomically since it will be
# called 4 times with the same fn
image.save(filename, format="PNG") # bits=16)
logger.debug("Saved image %s", filename)
image = None
except Exception:
logger.error("Cannot save image %s", filename)
raise
else:
buffered = io.BytesIO()
image.save(buffered, format="PNG", bits=16)
res = str(base64.b64encode(buffered.getvalue()), "ascii")
filename = "data:image/png;base64," + res
table = self.table
last = table.last()
if last is None or last["time"] != run_number:
values = {"filename": filename, "time": run_number}
table.add(values)
return self._return_run_step(self.state_blocked, steps_run=1)
def get_visualization(self) -> str:
return "heatmap"
def to_json(self, short: bool = False, with_speed: bool = True) -> JSon:
json = super(Heatmap, self).to_json(short, with_speed)
if short:
return json
return self.heatmap_to_json(json, short)
def heatmap_to_json(self, json: JSon, short: bool) -> JSon:
dfslot = self.get_input_slot("array")
assert isinstance(dfslot.output_module, Histogram2D)
histo: Histogram2D = dfslot.output_module
json["columns"] = [histo.x_column, histo.y_column]
histo_df = dfslot.data()
if histo_df is not None and len(histo_df) != 0:
row = histo_df.last()
if not (np.isnan(row["xmin"]) or np.isnan(row["xmax"])
or np.isnan(row["ymin"]) or np.isnan(row["ymax"])):
json["bounds"] = {
"xmin": row["xmin"],
"ymin": row["ymin"],
"xmax": row["xmax"],
"ymax": row["ymax"],
}
df = self.table
if df is not None and self._last_update != 0:
json["image"] = notNone(df.last())["filename"]
return json
def get_image(self, run_number: Optional[int] = None) -> Optional[str]:
filename: Optional[str]
table = self.table
if table is None or len(table) == 0:
return None
last = notNone(table.last())
# assert last is not None # len(table) > 0 so last is not None
if run_number is None or run_number >= last["time"]:
run_number = last["time"]
filename = last["filename"]
else:
time = table["time"]
idx = np.where(time == run_number)[0]
assert last is not None
if len(idx) == 0:
filename = last["filename"]
else:
filename = table["filename"][idx[0]]
return filename
def get_image_bin(self,
run_number: Optional[int] = None) -> Optional[bytes]:
file_url = self.get_image(run_number)
if file_url:
payload = file_url.split(",", 1)[1]
return base64.b64decode(payload)
else:
return None