Skip to content

⚡️ Speed up function make_norm_from_scale by 40% #244

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
codeflash-ai wants to merge 1 commit into
base: main
Choose a base branch
from
Open

⚡️ Speed up function make_norm_from_scale by 40% #244

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 42 additions & 17 deletions lib/matplotlib/colors.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1703,19 +1703,35 @@ class norm_cls(Normalize):
if base_norm_cls is None:
return functools.partial(make_norm_from_scale, scale_cls, init=init)

if isinstance(scale_cls, functools.partial):
# Avoid expensive isinstance check and attribute lookup: check type and attr directly
scale_args = ()
scale_kwargs_items = ()
if type(scale_cls) is functools.partial:
scale_args = scale_cls.args
scale_kwargs_items = tuple(scale_cls.keywords.items())
# In CPython, .keywords can be None or dict
scale_kwargs = scale_cls.keywords if scale_cls.keywords is not None else {}
scale_kwargs_items = tuple(scale_kwargs.items())
scale_cls = scale_cls.func
else:
scale_args = scale_kwargs_items = ()

# Cache the signature lookup for the default case
if init is None:
def init(vmin=None, vmax=None, clip=False): pass
# Use static signature, avoid function definition overhead
default_signature = inspect.Signature([
inspect.Parameter("vmin", inspect.Parameter.POSITIONAL_OR_KEYWORD, default=None),
inspect.Parameter("vmax", inspect.Parameter.POSITIONAL_OR_KEYWORD, default=None),
inspect.Parameter("clip", inspect.Parameter.POSITIONAL_OR_KEYWORD, default=False)
])
sig = default_signature
else:
# Signature is always reused within call
sig = inspect.signature(init)

# Fastest path to cached _make_norm_from_scale

return _make_norm_from_scale(
scale_cls, scale_args, scale_kwargs_items,
base_norm_cls, inspect.signature(init))
base_norm_cls, sig
)


@functools.cache
Expand Down Expand Up @@ -1746,8 +1762,9 @@ def __reduce__(self):
# class). If either import or attribute access fails, fall back to
# the general path.
try:
if cls is getattr(importlib.import_module(cls.__module__),
cls.__qualname__):
m = importlib.import_module(cls.__module__)
attr = getattr(m, cls.__qualname__)
if cls is attr:
return (_create_empty_object_of_class, (cls,), vars(self))
except (ImportError, AttributeError):
pass
Expand All @@ -1761,9 +1778,11 @@ def __init__(self, *args, **kwargs):
ba.apply_defaults()
super().__init__(
**{k: ba.arguments.pop(k) for k in ["vmin", "vmax", "clip"]})
# Avoid creating dict every call, build once at class construction
scale_kwargs = dict(scale_kwargs_items)
self._scale = functools.partial(
scale_cls, *scale_args, **dict(scale_kwargs_items))(
axis=None, **ba.arguments)
scale_cls, *scale_args, **scale_kwargs)(
axis=None, **ba.arguments)
self._trf = self._scale.get_transform()

__init__.__signature__ = bound_init_signature.replace(parameters=[
Expand All @@ -1772,18 +1791,22 @@ def __init__(self, *args, **kwargs):

def __call__(self, value, clip=None):
value, is_scalar = self.process_value(value)
if self.vmin is None or self.vmax is None:
vmin = self.vmin
vmax = self.vmax
if vmin is None or vmax is None:
self.autoscale_None(value)
if self.vmin > self.vmax:
vmin = self.vmin
vmax = self.vmax
if vmin > vmax:
raise ValueError("vmin must be less or equal to vmax")
if self.vmin == self.vmax:
if vmin == vmax:
return np.full_like(value, 0)
if clip is None:
clip = self.clip
if clip:
value = np.clip(value, self.vmin, self.vmax)
value = np.clip(value, vmin, vmax)
t_value = self._trf.transform(value).reshape(np.shape(value))
t_vmin, t_vmax = self._trf.transform([self.vmin, self.vmax])
t_vmin, t_vmax = self._trf.transform([vmin, vmax])
if not np.isfinite([t_vmin, t_vmax]).all():
raise ValueError("Invalid vmin or vmax")
t_value -= t_vmin
Expand All @@ -1794,9 +1817,11 @@ def __call__(self, value, clip=None):
def inverse(self, value):
if not self.scaled():
raise ValueError("Not invertible until scaled")
if self.vmin > self.vmax:
vmin = self.vmin
vmax = self.vmax
if vmin > vmax:
raise ValueError("vmin must be less or equal to vmax")
t_vmin, t_vmax = self._trf.transform([self.vmin, self.vmax])
t_vmin, t_vmax = self._trf.transform([vmin, vmax])
if not np.isfinite([t_vmin, t_vmax]).all():
raise ValueError("Invalid vmin or vmax")
value, is_scalar = self.process_value(value)
Expand Down