add generated ANCOVA dataset + added to tests

Author: drbenvincent

causalpy/data/ancova_generated.csv

@@ -0,0 +1,201 @@
+group,pre,post
+1,12.264858355350702,14.642281485826853
+0,9.927449994409082,10.247109509774678
+1,11.143182680959526,13.208891968269526
+1,11.581027267217818,13.176653832648514
+0,10.212011893570054,9.588953081804807
+0,10.099040162649114,10.485018086502661
+0,9.848261857556366,10.17443467781112
+0,11.248891593197538,11.44914346555337
+0,9.988520810244031,9.982947871444773
+1,11.617199990961508,13.968112204434641
+1,11.357258461761651,13.086386661748689
+1,11.443642992680244,14.117123920581372
+1,12.418698299591519,14.556239848432652
+1,11.08695058268037,13.75020247980409
+1,11.903073922037589,14.37467647038186
+0,8.259941992017877,8.24264463958351
+1,11.260206775093495,13.38839323219005
+0,8.84328240290688,8.991555943487812
+1,12.833944746113795,14.237136574260715
+0,8.243694198122864,8.167157809735029
+0,9.549873086453907,9.482137435602832
+0,7.835284682165309,7.612336136673208
+1,12.88752541631879,14.20156778126195
+0,9.292511464335078,9.523798988308581
+0,6.719488929307644,6.828074824534081
+0,12.034900617958858,11.972499594137854
+1,12.8017475187411,15.201569801046992
+1,11.192165456103698,12.538373263956393
+0,9.218026719316452,9.030263180794165
+0,9.794321887202216,10.198486677176872
+1,11.20016101928247,12.824166571584524
+1,11.810122517971482,14.126998873357666
+1,12.113003771192073,14.036824795467137
+0,9.853702428392674,10.134049576701152
+1,11.581812146609074,13.07485589712462
+0,11.44840652163122,11.69112926558477
+0,9.349158155658037,9.500824731689992
+1,12.307317701635883,14.032196697704606
+1,11.521163845772623,13.538391033781803
+0,10.009662624335466,9.887190209442021
+0,9.033281858777084,9.338886734860036
+1,11.807639539107594,13.685887202945342
+1,12.261871954119055,14.28611940629696
+1,11.689690609394821,13.76124483586917
+0,8.911128292748554,8.993488979563084
+1,13.986389194790188,15.995992132088245
+1,11.45153204238965,14.049822896509092
+0,10.627081564978075,11.355650937465569
+1,11.480296147929169,13.20034610034534
+0,10.58098947444771,10.206358987350022
+1,11.648239555454863,14.089331157975257
+0,10.981675527174431,11.055058766740057
+0,10.50965455054716,11.04747252077189
+0,10.625315126287015,10.318955841019823
+1,11.554557224238774,12.85449149777984
+0,10.900812128840057,11.6259543261164
+1,11.854010777811842,13.489847911072802
+1,12.817650208204363,14.870388232209867
+1,10.848033033225063,12.839913751596841
+0,10.966934906053474,10.813747590173843
+0,10.520610302205263,10.796120263549959
+0,10.809431727759202,11.12251292573233
+1,12.317390921086742,14.945538687160642
+0,9.030030538887935,9.561708783696691
+0,9.693392158056453,10.018650351124018
+0,9.76386068297627,9.688878210224578
+0,9.467777687386512,9.344090442163484
+1,10.904517888189673,12.297713114527422
+1,11.905084055656603,13.759863062869309
+0,6.369206844201461,5.77845021451916
+0,8.894134569538416,8.254637265417841
+0,9.594751482931796,10.118800810581622
+1,12.058985852615915,14.667351520008605
+0,10.238595428637137,10.791232399316057
+1,12.158784012936945,14.054015730061467
+0,9.814345081491227,10.601782581217433
+1,12.141122803262492,14.656980603576939
+0,9.227043335151656,9.426941715795842
+0,9.903178883200862,8.765326384529072
+0,10.126966921640609,9.489665602782823
+1,11.13588268952632,12.876322536297383
+1,11.300107989923264,13.327382871088595
+1,13.272978942338987,15.585840771500374
+0,9.14906665334288,9.296816374189518
+1,13.393604977690963,14.675487876788951
+0,9.192828412985934,9.379279238594563
+0,9.373788708805776,9.478494343512358
+0,10.373292219777179,10.559437714400707
+1,11.961585386183025,13.833246351271987
+1,11.570707466740528,13.680172672911127
+1,13.567812559913143,15.135868319389388
+0,9.786454645494572,9.580975708036352
+1,12.212608107180793,13.694232107313065
+0,9.397077261006636,9.401053454968341
+1,10.989718910446546,12.231612211172177
+0,10.179643430531856,10.187005861094837
+0,10.661614250233816,10.450741283835002
+1,12.517893975625041,14.242553055257599
+0,8.186534061787329,7.769197674208002
+0,8.430748557999433,8.105732205758365
+0,9.948912626255728,9.990095979646151
+1,11.119208147257408,13.448160765847305
+0,10.116844072623824,10.044463935297504
+1,13.20734255110846,14.715956833385809
+0,10.058816735912718,9.568920513635073
+1,13.399131461277102,15.898598531725778
+0,9.183861278717954,9.05817250218202
+1,12.93215649379921,15.201097651149194
+1,11.247780032965364,13.31728117568601
+0,9.474570753816382,9.321507957947217
+0,11.029863789865384,10.617870761906703
+1,10.612534743385979,12.21293067230243
+0,8.162204171003342,8.161557410890328
+0,8.364184228934858,8.027395466630457
+1,13.231498472533868,15.460289619087318
+1,11.495426056170695,13.389206276180522
+1,13.597246796764637,16.25253360430974
+0,10.715558413738783,11.292696744480372
+1,11.373960868330322,12.840506571893563
+1,11.860168962863618,13.23180940040767
+1,12.155085338079964,14.011800436576367
+1,12.461259000750996,13.613506772155183
+0,10.640933127724542,10.775297359456811
+0,9.667968556475836,9.50939678375694
+0,10.294149423598189,10.196766564686822
+0,11.933900610983688,12.086776790684851
+0,9.872039239445167,10.160105674157954
+1,13.033479847041354,14.854905709755855
+1,12.026629349841249,13.552620635900984
+1,11.879439414827207,13.150523728615966
+0,9.566686505463542,9.704222475950209
+0,10.854664037419726,11.38175393879568
+1,11.098046113824383,12.62787100050706
+0,9.782165369675669,9.499915875474162
+1,10.946363184143978,12.839913823256849
+1,12.499486260308794,14.659496064291226
+1,10.993681063536865,12.232856790327482
+1,12.14072997893197,13.520399299749684
+1,12.741360844286929,14.093733041578513
+0,9.211192222054821,9.383914042257693
+1,13.913764933576681,15.90230836727072
+0,9.549618940826129,9.550021154993763
+0,10.900557041510229,11.004982845659425
+1,10.78976082147783,12.258212585685055
+0,9.610091590761158,9.288806688882772
+1,12.87082164331867,14.855594819083063
+1,13.226763655568835,14.741675343660166
+1,10.509134781612119,12.278683775650363
+0,9.392946687346507,9.251414329990293
+1,11.654899514693925,13.66826270458631
+1,12.345996772351288,13.908552145921075
+0,8.936573461339567,8.263868160092676
+1,13.444517967913384,15.155337552450447
+0,8.88648505864527,9.277111903412615
+1,11.143959465125791,13.464548955240314
+1,12.440516168823322,13.752017958160309
+0,10.56959536418749,9.993594680331267
+1,12.362879640914342,14.425151735546999
+1,14.754321117020414,16.66185852912643
+0,8.748985743665008,9.090324460954626
+1,12.116637901385804,13.890990060165413
+0,11.310129837671347,11.815079555186793
+0,9.056724450312705,8.908699440692132
+1,10.6713558110713,13.27267710957835
+1,11.768996408549217,13.609890869812782
+1,10.987663898349393,12.959770406961267
+0,9.271347690394514,8.972284210261076
+1,13.41062647812046,15.643150501577235
+1,12.080925001031122,13.969090283189457
+1,12.131309628632286,14.081290096923862
+0,9.148662952434565,9.959056582869447
+1,11.659141704267826,13.288607339068161
+1,13.730246049103854,16.332346664778708
+0,10.493705379414665,10.29265378097157
+0,9.61510292897469,10.005525775760427
+1,12.097005139591186,13.55533637332416
+0,9.70061849575247,10.161630504822973
+0,9.457096609096899,10.010365868615919
+0,10.677322714363163,10.80760495984205
+1,10.513325497923493,12.126984080601527
+1,12.062551901708142,14.712645533163277
+1,13.33252423398856,14.837332146889565
+1,13.953036672884672,16.20436478075008
+1,12.343834047799021,14.278233304758707
+0,9.730658103486451,9.471795169542592
+0,11.980654709795681,11.831425438132943
+1,12.663214987320673,14.739223146205598
+0,9.489950507583748,8.893683024254251
+1,10.33133189064659,12.532810875624302
+0,11.691805096688345,11.945931314422928
+1,11.714517371918909,13.468702907081228
+1,11.96599263232672,13.936729757214527
+1,14.185092904808357,15.587409097367365
+0,9.782312952391104,9.299757499093808
+0,10.79382217804355,11.450147359966014
+0,9.939731553868997,9.876398697296123
+0,9.064563009578302,9.19012720728962
+0,12.342369459765722,12.215277731909632
+1,11.721192710169209,13.235723336589501
+0,10.5218333495686,10.177956671617448

causalpy/data/datasets.py

@@ -14,6 +14,7 @@
     "its simple": {"filename": "its_simple.csv"},
     "rd": {"filename": "regression_discontinuity.csv"},
     "sc": {"filename": "synthetic_control.csv"},
+    "anova1": {"filename": "ancova_generated.csv"},
 }
 
 

causalpy/data/simulate_data.py

@@ -210,3 +210,13 @@ def impact(x):
     y = np.sin(x * 3) + impact(x) + norm.rvs(scale=0.1, size=N)
 
     return pd.DataFrame({"x": x, "y": y, "treated": is_treated(x)})
+
+
+def generate_ancova_data(
+    N=200, pre_treatment_means=np.array([10, 12]), treatment_effect=2
+):
+    group = np.random.choice(2, size=N)
+    pre = np.random.normal(loc=pre_treatment_means[group])
+    post = pre + treatment_effect * group + np.random.normal(size=N) * 0.4
+    df = pd.DataFrame({"group": group, "pre": pre, "post": post})
+    return df

causalpy/tests/test_data_loading.py

@@ -3,7 +3,18 @@
 
 import causalpy as cp
 
-tests = ["banks", "brexit", "covid", "did", "drinking", "its", "its simple", "rd", "sc"]
+tests = [
+    "banks",
+    "brexit",
+    "covid",
+    "did",
+    "drinking",
+    "its",
+    "its simple",
+    "rd",
+    "sc",
+    "anova1",
+]
 
 
 @pytest.mark.parametrize("dataset_name", tests)