plotly.graph_objects.FigureWidget does not support NaN or Inf due to JSON limitations

Hello everyone.

I'm trying to plot some interactive heatmaps on a Jupyter Notebook using Ipywidgets and FigureWidget.
These heatmaps, by design, will contain some NaN values that needs to be "not plotted" as such. From this point of view, my previous experience of go.Figure was perfect.

However, when trying to use FigureWidget, I get the ValueError: Out of range float values are not JSON compliant, as JSON base standard is not able to serialize Infs or NaNs. Is it possible to avoid this limitation somehow?

A minimal example to execute in a notebook:

# %%
import numpy as np
import plotly.graph_objects as go


# %%
matrix_with_no_nans = np.random.uniform(0, 1, size=(100, 100))

matrix_with_some_nans = matrix_with_no_nans.copy()
matrix_with_some_nans[matrix_with_some_nans < 0.5] = np.nan


# %%
fig_1 = go.FigureWidget(
    data=go.Heatmap(z=matrix_with_no_nans)
)
# Works as intended
fig_1


# %%
fig_2 = go.FigureWidget(
    data=go.Heatmap(z=matrix_with_some_nans)
)
# Json error!
fig_2

# %%

The full text of the error message from my personal notebook:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/tmp/ipykernel_129118/716380412.py in <module>
----> 1 fig_2 = go.FigureWidget(
      2     data=go.Heatmap(z=matrix_with_some_nans)
      3 )
      4 fig_2

~/anaconda3/lib/python3.8/site-packages/plotly/graph_objs/_figurewidget.py in __init__(self, data, layout, frames, skip_invalid, **kwargs)
    588             is invalid AND skip_invalid is False
    589         """
--> 590         super(FigureWidget, self).__init__(data, layout, frames, skip_invalid, **kwargs)
    591 
    592     def add_bar(

~/anaconda3/lib/python3.8/site-packages/plotly/basewidget.py in __init__(self, data, layout, frames, skip_invalid, **kwargs)
    115         # with the `layout` constructor parameter of the `widgets.DOMWidget`
    116         # ipywidgets class
--> 117         super(BaseFigureWidget, self).__init__(
    118             data=data,
    119             layout_plotly=layout,

~/anaconda3/lib/python3.8/site-packages/plotly/basedatatypes.py in __init__(self, data, layout_plotly, frames, skip_invalid, **kwargs)
    524         # The _data property is a list of dicts containing the properties
    525         # explicitly set by the user for each trace.
--> 526         self._data = [deepcopy(trace._props) for trace in data]
    527 
    528         # ### Create data defaults ###

~/anaconda3/lib/python3.8/site-packages/plotly/basedatatypes.py in __setattr__(self, prop, value)
    719         if prop.startswith("_") or hasattr(self, prop):
    720             # Let known properties and private properties through
--> 721             super(BaseFigure, self).__setattr__(prop, value)
    722         else:
    723             # Raise error on unknown public properties

~/anaconda3/lib/python3.8/site-packages/traitlets/traitlets.py in __set__(self, obj, value)
    604             raise TraitError('The "%s" trait is read-only.' % self.name)
    605         else:
--> 606             self.set(obj, value)
    607 
    608     def _validate(self, obj, value):

~/anaconda3/lib/python3.8/site-packages/traitlets/traitlets.py in set(self, obj, value)
   2649             return super().set(obj, [value])
   2650         else:
-> 2651             return super().set(obj, value)
   2652 
   2653 

~/anaconda3/lib/python3.8/site-packages/traitlets/traitlets.py in set(self, obj, value)
    593             # we explicitly compare silent to True just in case the equality
    594             # comparison above returns something other than True/False
--> 595             obj._notify_trait(self.name, old_value, new_value)
    596 
    597     def __set__(self, obj, value):

~/anaconda3/lib/python3.8/site-packages/traitlets/traitlets.py in _notify_trait(self, name, old_value, new_value)
   1217 
   1218     def _notify_trait(self, name, old_value, new_value):
-> 1219         self.notify_change(Bunch(
   1220             name=name,
   1221             old=old_value,

~/anaconda3/lib/python3.8/site-packages/ipywidgets/widgets/widget.py in notify_change(self, change)
    603             if name in self.keys and self._should_send_property(name, getattr(self, name)):
    604                 # Send new state to front-end
--> 605                 self.send_state(key=name)
    606         super(Widget, self).notify_change(change)
    607 

~/anaconda3/lib/python3.8/site-packages/ipywidgets/widgets/widget.py in send_state(self, key)
    487             state, buffer_paths, buffers = _remove_buffers(state)
    488             msg = {'method': 'update', 'state': state, 'buffer_paths': buffer_paths}
--> 489             self._send(msg, buffers=buffers)
    490 
    491 

~/anaconda3/lib/python3.8/site-packages/ipywidgets/widgets/widget.py in _send(self, msg, buffers)
    735         """Sends a message to the model in the front-end."""
    736         if self.comm is not None and self.comm.kernel is not None:
--> 737             self.comm.send(data=msg, buffers=buffers)
    738 
    739     def _repr_keys(self):

~/anaconda3/lib/python3.8/site-packages/ipykernel/comm/comm.py in send(self, data, metadata, buffers)
    120     def send(self, data=None, metadata=None, buffers=None):
    121         """Send a message to the frontend-side version of this comm"""
--> 122         self._publish_msg('comm_msg',
    123             data=data, metadata=metadata, buffers=buffers,
    124         )

~/anaconda3/lib/python3.8/site-packages/ipykernel/comm/comm.py in _publish_msg(self, msg_type, data, metadata, buffers, **keys)
     64         metadata = {} if metadata is None else metadata
     65         content = json_clean(dict(data=data, comm_id=self.comm_id, **keys))
---> 66         self.kernel.session.send(self.kernel.iopub_socket, msg_type,
     67             content,
     68             metadata=json_clean(metadata),

~/anaconda3/lib/python3.8/site-packages/jupyter_client/session.py in send(self, stream, msg_or_type, content, parent, ident, buffers, track, header, metadata)
    828         if self.adapt_version:
    829             msg = adapt(msg, self.adapt_version)
--> 830         to_send = self.serialize(msg, ident)
    831         to_send.extend(buffers)
    832         longest = max([len(s) for s in to_send])

~/anaconda3/lib/python3.8/site-packages/jupyter_client/session.py in serialize(self, msg, ident)
    702             content = self.none
    703         elif isinstance(content, dict):
--> 704             content = self.pack(content)
    705         elif isinstance(content, bytes):
    706             # content is already packed, as in a relayed message

~/anaconda3/lib/python3.8/site-packages/jupyter_client/session.py in json_packer(obj)
     93 
     94 def json_packer(obj):
---> 95     return jsonapi.dumps(
     96         obj,
     97         default=json_default,

~/anaconda3/lib/python3.8/site-packages/zmq/utils/jsonapi.py in dumps(o, **kwargs)
     23     Keyword arguments are passed along to :py:func:`json.dumps`.
     24     """
---> 25     return json.dumps(o, **kwargs).encode("utf8")
     26 
     27 

~/anaconda3/lib/python3.8/json/__init__.py in dumps(obj, skipkeys, ensure_ascii, check_circular, allow_nan, cls, indent, separators, default, sort_keys, **kw)
    232     if cls is None:
    233         cls = JSONEncoder
--> 234     return cls(
    235         skipkeys=skipkeys, ensure_ascii=ensure_ascii,
    236         check_circular=check_circular, allow_nan=allow_nan, indent=indent,

~/anaconda3/lib/python3.8/json/encoder.py in encode(self, o)
    197         # exceptions aren't as detailed.  The list call should be roughly
    198         # equivalent to the PySequence_Fast that ''.join() would do.
--> 199         chunks = self.iterencode(o, _one_shot=True)
    200         if not isinstance(chunks, (list, tuple)):
    201             chunks = list(chunks)

~/anaconda3/lib/python3.8/json/encoder.py in iterencode(self, o, _one_shot)
    255                 self.key_separator, self.item_separator, self.sort_keys,
    256                 self.skipkeys, _one_shot)
--> 257         return _iterencode(o, 0)
    258 
    259 def _make_iterencode(markers, _default, _encoder, _indent, _floatstr,

ValueError: Out of range float values are not JSON compliant

Thanks!

Comienza reproduciendo el ejemplo mínimo de mapa de calor de FigureWidget y sigue FigureWidget a través de plotly/basewidget.py hasta la ruta de ipywidgets send_state mostrada en el traceback. Se considera terminado cuando FigureWidget acepta los datos del mapa de calor con NaN/Inf sin producir el JSON ValueError y conserva las celdas que intencionadamente no se representan.