[docs]defget_cost(name:str,*,backend:str="moscot",**kwargs:Any)->Any:"""Get cost function for a specific backend."""key=f"{backend}{_SEP}{name}"ifkeynotin_REGISTRY:raiseValueError(f"Cost `{name!r}` is not available for backend `{backend!r}`.")return_REGISTRY[key](**kwargs)
[docs]defget_available_costs(backend:Optional[str]=None)->Dict[str,Tuple[str,...]]:"""Return available costs. Parameters ---------- backend Select cost specific to a backend. If :obj:`None`, return the costs for each backend. Returns ------- Dictionary with keys as backend names and values as registered cost functions. """groups:Dict[str,List[str]]=_get_available_backends_and_costs()ifbackendisNone:return{k:tuple(v)fork,vingroups.items()}ifbackendnotingroups:raiseKeyError(f"No backend named `{backend!r}`.")return{backend:tuple(groups[backend])}
[docs]defregister_cost(name:str,*,backend:str)->Any:"""Register cost function for a specific backend."""return_REGISTRY.register(f"{backend}{_SEP}{name}")
def_get_available_backends_and_costs():"""Return a dictionary of available backends with their corresponding list of costs. Returns ------- Default dictionary with keys as backend names and values as registered cost functions. """groups:Dict[str,List[str]]=collections.defaultdict(list)forkeyin_REGISTRY:back,*name=key.split(_SEP)groups[back].append(_SEP.join(name))returngroups