Dynamic limits for element-wise checks

[oniht]

I've been going through the documentation and trying to figure out if it's possible to create an element-wise greater than check where the limit is different for each group defined by another column. Ideally, the limit would be inferred from the data. (e.g. median - 3 times standard deviation for that group)

So far I haven't found out any solution. I would appreciate help if it's indeed possible to such a check.

[jonathanstathakis]

Hi oniht, I am by no means an expert in Pandera, but if I correctly understand your problem, then I may have a solution via a custom check. I am using the class based API but I believe you can achieve this through the object API if needed. Please see the code block below:

import pandera.extensions as extensions
import pandas as pd
import pandera as pa

# define an example df
df = pd.DataFrame(
    {
    "x": [1, 2, 3, 10],
    "y": [-3, 4, 5, 7],
    "label": ["b", "a", "b", "a"]
     }
    )


@extensions.register_check_method(statistics=["check_key", "grouping_key"])
def check_groups_are_within_3_std(df, *, check_key: str, grouping_key: str):
    """
    A custom check testing whether each groups median is less than or equal to their
    3 sigma
    """
    check = df.groupby([grouping_key])[check_key].apply(
        lambda x: x.median() <= x.std() * 3
    )

    return check


class Schema(pa.DataFrameModel):

    x: int
    y: int
    label: str

    class Config:
        # register the custom check
        check_groups_are_within_3_std = {"check_key": "x", "grouping_key": "label"}

Let me know how that goes, Im happy to follow up.

[cosmicBboy]

@oniht based on the context provided this approach should work

[oniht]

Hi @jonathanstathakis and @cosmicBboy , thank you for the quick reply! This is close but not entirely what I'm looking for. I would like to have an element-wise check such that
e >= x.median() - x.std() * 3 however I think I now find a working solution with df.groupby(grouping_key, group_keys=False)[check_key].apply(lambda x: x >= (x.median() - x.std() * 3)) in the custom check. Could this work?

[cosmicBboy]

seems like that should work, can you give it a try?