# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
# %load_ext autoreload
# %autoreload 2
# %aimport utils_1_1

import pandas as pd
import numpy as np
import altair as alt
from altair_saver import save
import datetime
import dateutil.parser
from os.path import join

from constants_1_1 import SITE_FILE_TYPES
from utils_1_1 import (
    get_site_file_paths,
    get_site_file_info,
    get_site_ids,
    get_visualization_subtitle,
    get_country_color_map,
    apply_theme,
)
from web import for_website

alt.data_transformers.disable_max_rows(); # Allow using rows more than 5000
# -

# > I just uploaded charlson results for post admission too by country. The file that was originally there and that you had previously visualized is for pre-admission. I’ve renamed the files appropriately. Could you revisualize to show charlson trends both pre admission and post admission b/w wave by country?

# +
df = pd.read_csv(join("..", "data", "Phase2.1", "temporal", "charlson", "charlson_countrylevel_preadmission.csv"))

df = df.loc[~df["country"].isin(["SINGAPORE"])]
df["country"] = df["country"].apply(lambda name: name if name in ["USA", "UK"] else name[0] + name.lower()[1:])
df["setting"] = df["setting"].apply(lambda s: s[0].upper() + s[1:])

orig_df = df

cs_df = pd.DataFrame()
for (country, setting), country_setting_df in df.groupby(by=["country", "setting"]):    
    cs_df = cs_df.append({
        "country": country,
        "setting": setting,
        "charlson_early": country_setting_df.loc[country_setting_df["wave"] == "early"].iloc[0]["charlson"],
        "charlson_late": country_setting_df.loc[country_setting_df["wave"] == "late"].iloc[0]["charlson"]
    }, ignore_index=True)
df = cs_df
df["difference"] = df["charlson_late"] - df["charlson_early"]
df["is_left_direction"] = df["difference"] < 0.0
df["is_right_direction"] = df["difference"] > 0.0
df["min"] = df.apply(lambda row: min(row["charlson_late"], row["charlson_early"]), axis='columns')
df.head()
# -

COUNTRY_COLORS = ['#0072B2', '#E79F00', '#029F73', '#D45E00', '#CB7AA7']
COUNTRY_NAMES = []

# +
bars = alt.Chart().mark_bar(size=10).encode(
    x=alt.X("charlson_early:Q", axis=alt.Axis(title="Charlson index")),
    x2=alt.X2("charlson_late:Q"),
    y=alt.Y("setting:N", axis=alt.Axis(title="Severity")),
    color=alt.Color("setting:N", legend=alt.Legend(title="Severity"))
)

triangle_opacity_scale = scale=alt.Scale(domain=[True, False], range=[1, 0])

triangles_left = alt.Chart().mark_point(shape="triangle-left", filled=True, size=100, xOffset=-4, opacity=1).encode(
    x=alt.X("charlson_late:Q", axis=alt.Axis(title="Charlson index")),
    y=alt.Y("setting:N", axis=alt.Axis(title="Severity")),
    color=alt.Color("setting:N", legend=alt.Legend(title="Severity")),
    opacity=alt.Opacity("is_left_direction:Q", scale=triangle_opacity_scale, legend=None)
)

triangles_right = alt.Chart().mark_point(shape="triangle-right", filled=True, size=100, xOffset=8, opacity=1).encode(
    x=alt.X("charlson_early:Q", axis=alt.Axis(title="Charlson index")),
    y=alt.Y("setting:N", axis=alt.Axis(title="Severity")),
    color=alt.Color("setting:N", legend=alt.Legend(title="Severity")),
    opacity=alt.Opacity("is_right_direction:Q", scale=triangle_opacity_scale, legend=None)
)

chart = alt.layer(bars, triangles_left, triangles_right, data=df).facet(
    row=alt.Row("country:N", header=alt.Header(title="Country")),
)

chart = chart.properties(
    title={
        "text": "Charlson index from early wave to late wave",
        "subtitle": "Severity: All | Status: Pre-admission",
        "subtitleColor": "gray",
        "dx": 130
    },
)

chart = apply_theme(chart, title_anchor='start')

chart

# +
left_bars = alt.Chart().mark_bar(size=14, cornerRadiusTopLeft=9, cornerRadiusBottomLeft=9).encode(
    x=alt.X("charlson_early:Q", axis=alt.Axis(title="Charlson index")),
    x2=alt.X2("charlson_late:Q"),
    y=alt.Y("setting:N", axis=alt.Axis(title="Severity")),
    color=alt.Color("setting:N", legend=alt.Legend(title="Severity")),
    opacity=alt.Opacity("is_left_direction:Q", scale=triangle_opacity_scale, legend=None)
).properties(height=90)

right_bars = alt.Chart().mark_bar(size=14, cornerRadiusTopRight=9, cornerRadiusBottomRight=9).encode(
    x=alt.X("charlson_early:Q", axis=alt.Axis(title="Charlson index")),
    x2=alt.X2("charlson_late:Q"),
    y=alt.Y("setting:N", axis=alt.Axis(title="Severity")),
    color=alt.Color("setting:N", legend=alt.Legend(title="Severity")),
    opacity=alt.Opacity("is_right_direction:Q", scale=triangle_opacity_scale, legend=None)
)

chart = alt.layer(left_bars, right_bars, data=df).facet(
    row=alt.Row("country:N", header=alt.Header(title="Country")),
)

chart = chart.properties(
    title={
        "text": "Charlson index from early wave to late wave",
        "subtitle": "Severity: All | Status: Pre-admission",
        "subtitleColor": "gray",
        "dx": 130
    },
)

chart = apply_theme(chart, title_anchor='start')

chart

# +
all_df = df.loc[df["setting"] == "All"]

index_domain = [all_df[["charlson_early", "charlson_late"]].min().min() - 1, all_df[["charlson_early", "charlson_late"]].max().max() + 1]

# Sort countries from smallest shift to largest shift

country_shift_sort = all_df.sort_values(by="difference", ascending=False)["country"].tolist()

bars = alt.Chart().mark_bar(size=18).encode(
    x=alt.X("charlson_early:Q", scale=alt.Scale(domain=index_domain), axis=alt.Axis(title="Charlson index")),
    x2=alt.X2("charlson_late:Q"),
    y=alt.Y("country:N", axis=alt.Axis(title="Country"), scale=alt.Scale(domain=country_shift_sort)),
    color=alt.Color("country:N", legend=alt.Legend(title="Country")),
)

triangle_opacity_scale = scale=alt.Scale(domain=[True, False], range=[1, 0])

triangles_left = alt.Chart().mark_point(shape="triangle-left", filled=True, size=300, xOffset=-7, opacity=1).encode(
    x=alt.X("charlson_late:Q", axis=alt.Axis(title="Charlson index")),
    y=alt.Y("country:N", axis=alt.Axis(title="Country")),
    color=alt.Color("country:N", legend=alt.Legend(title="Country")),
    opacity=alt.Opacity("is_left_direction:Q", scale=triangle_opacity_scale, legend=None)
)

triangles_right = alt.Chart().mark_point(shape="triangle-right", filled=True, size=300, xOffset=8, opacity=1).encode(
    x=alt.X("charlson_early:Q", axis=alt.Axis(title="Charlson index")),
    y=alt.Y("country:N", axis=alt.Axis(title="Country")),
    color=alt.Color("country:N", legend=alt.Legend(title="Country")),
    opacity=alt.Opacity("is_right_direction:Q", scale=triangle_opacity_scale, legend=None)
)

chart = alt.layer(bars, triangles_left, triangles_right, data=all_df)

chart = chart.properties(
    title={
        "text": "Charlson index from early wave to late wave",
        "subtitle": "Severity: All | Status: Pre-admission",
        "subtitleColor": "gray",
        "dx": 91
    },
    height=200
)

chart = apply_theme(chart, title_anchor='start')

chart

# +
all_df = df.loc[df["setting"] == "All"]

index_domain = [all_df[["charlson_early", "charlson_late"]].min().min() - 1, all_df[["charlson_early", "charlson_late"]].max().max() + 1]

# Sort countries from smallest shift to largest shift

country_shift_sort = all_df.sort_values(by="difference", ascending=False)["country"].tolist()

bars = alt.Chart().mark_bar(size=18, cornerRadiusTopLeft=9, cornerRadiusBottomLeft=9).encode(
    x=alt.X("charlson_early:Q", scale=alt.Scale(domain=index_domain), axis=alt.Axis(title="Charlson index")),
    x2=alt.X2("charlson_late:Q"),
    y=alt.Y("country:N", axis=alt.Axis(title="Country"), scale=alt.Scale(domain=country_shift_sort)),
    color=alt.Color("country:N", legend=alt.Legend(title="Country")),
)

chart = alt.layer(bars, data=all_df)

chart = chart.properties(
    title={
        "text": "Charlson index from early wave to late wave",
        "subtitle": "Severity: All | Status: Pre-admission",
        "subtitleColor": "gray",
        "dx": 91
    },
    height=200
)

chart = apply_theme(chart, title_anchor='start')

chart

# +
all_df = df.loc[df["setting"] == "All"]

index_domain = [all_df[["charlson_early", "charlson_late"]].min().min() - 1, all_df[["charlson_early", "charlson_late"]].max().max() + 1]

# Sort countries by their late wave index value, ascending

country_shift_sort = all_df.sort_values(by="charlson_late", ascending=True)["country"].tolist()

bars = alt.Chart().mark_bar(size=18, cornerRadiusTopLeft=9, cornerRadiusBottomLeft=9).encode(
    x=alt.X("charlson_early:Q", scale=alt.Scale(domain=index_domain), axis=alt.Axis(title="Charlson index")),
    x2=alt.X2("charlson_late:Q"),
    y=alt.Y("country:N", axis=alt.Axis(title="Country"), scale=alt.Scale(domain=country_shift_sort)),
    color=alt.Color("country:N", legend=alt.Legend(title="Country")),
)

chart = alt.layer(bars, data=all_df)

chart = chart.properties(
    title={
        "text": "Charlson index from early wave to late wave",
        "subtitle": "Severity: All | Status: Pre-admission",
        "subtitleColor": "gray",
        "dx": 91
    },
    height=200
)

chart = apply_theme(chart, title_anchor='start')

chart

# +
all_df = df.loc[df["setting"] == "All"]

index_domain = [all_df[["charlson_early", "charlson_late"]].min().min() - 1, all_df[["charlson_early", "charlson_late"]].max().max() + 1]

# Sort countries by their early wave index value, ascending

country_shift_sort = all_df.sort_values(by="charlson_early", ascending=True)["country"].tolist()

bars = alt.Chart().mark_bar(size=18, cornerRadiusTopLeft=9, cornerRadiusBottomLeft=9).encode(
    x=alt.X("charlson_early:Q", scale=alt.Scale(domain=index_domain), axis=alt.Axis(title="Charlson index")),
    x2=alt.X2("charlson_late:Q"),
    y=alt.Y("country:N", axis=alt.Axis(title="Country"), scale=alt.Scale(domain=country_shift_sort)),
    color=alt.Color("country:N", legend=alt.Legend(title="Country")),
)

chart = alt.layer(bars, data=all_df)

chart = chart.properties(
    title={
        "text": "Charlson index from early wave to late wave",
        "subtitle": "Severity: All | Status: Pre-admission",
        "subtitleColor": "gray",
        "dx": 91
    },
    height=200
)

chart = apply_theme(chart, title_anchor='start')

chart

# +
severe_df = df.loc[df["setting"] == "Ever"]

index_domain = [severe_df[["charlson_early", "charlson_late"]].min().min() - 1, severe_df[["charlson_early", "charlson_late"]].max().max() + 1]

# Sort countries by their early wave index value, ascending

country_shift_sort = severe_df.sort_values(by="charlson_early", ascending=True)["country"].tolist()

left_bars = alt.Chart().mark_bar(size=18, cornerRadiusTopLeft=9, cornerRadiusBottomLeft=9).encode(
    x=alt.X("charlson_early:Q", scale=alt.Scale(domain=index_domain), axis=alt.Axis(title="Charlson index")),
    x2=alt.X2("charlson_late:Q"),
    y=alt.Y("country:N", axis=alt.Axis(title="Country"), scale=alt.Scale(domain=country_shift_sort)),
    color=alt.Color("country:N", legend=alt.Legend(title="Country")),
    opacity=alt.Opacity("is_left_direction:Q", scale=triangle_opacity_scale, legend=None)
)

right_bars = alt.Chart().mark_bar(size=18, cornerRadiusTopRight=9, cornerRadiusBottomRight=9).encode(
    x=alt.X("charlson_early:Q", scale=alt.Scale(domain=index_domain), axis=alt.Axis(title="Charlson index")),
    x2=alt.X2("charlson_late:Q"),
    y=alt.Y("country:N", axis=alt.Axis(title="Country"), scale=alt.Scale(domain=country_shift_sort)),
    color=alt.Color("country:N", legend=alt.Legend(title="Country")),
    opacity=alt.Opacity("is_right_direction:Q", scale=triangle_opacity_scale, legend=None)
)

chart = alt.layer(left_bars, right_bars, data=severe_df)

chart = chart.properties(
    title={
        "text": "Charlson index from early wave to late wave",
        "subtitle": "Severity: Ever | Status: Pre-admission",
        "subtitleColor": "gray",
        "dx": 91,
    },
    height=200
)

chart = apply_theme(chart, title_anchor='start')

chart

# +

index_domain = [df[["charlson_early", "charlson_late"]].min().min() - 1, df[["charlson_early", "charlson_late"]].max().max() + 1]

# Sort countries by their early wave index value, ascending

country_shift_sort = df.sort_values(by="charlson_early", ascending=True)["country"].tolist()

bars = alt.Chart(df).mark_bar(size=18, cornerRadiusTopLeft=9, cornerRadiusBottomLeft=9).encode(
    x=alt.X("charlson_early:Q", scale=alt.Scale(domain=index_domain), axis=alt.Axis(title="Charlson index")),
    x2=alt.X2("charlson_late:Q"),
    y=alt.Y("country:N", axis=alt.Axis(title="Country"), scale=alt.Scale(domain=country_shift_sort)),
    color=alt.Color("country:N", legend=alt.Legend(title="Country")),
    row=alt.Row("setting:N", header=alt.Header(title="Severity"))
)

chart = bars.properties(
    title={
        "text": "Charlson index from early wave to late wave",
        "subtitle": "Status: Pre-admission",
        "subtitleColor": "gray",
        "dx": 150
    },
    height=140
)

chart = apply_theme(chart, title_anchor='start')

chart

# +
left_bars = alt.Chart().mark_bar(size=14, cornerRadiusTopLeft=9, cornerRadiusBottomLeft=9).encode(
    x=alt.X("charlson_early:Q", axis=alt.Axis(title="Charlson index")),
    x2=alt.X2("charlson_late:Q"),
    y=alt.Y("country:N", axis=alt.Axis(title="Country"), scale=alt.Scale(domain=country_shift_sort)),
    color=alt.Color("country:N", legend=alt.Legend(title="Country")),
    opacity=alt.Opacity("is_left_direction:Q", scale=triangle_opacity_scale, legend=None)
).properties(height=140)

right_bars = alt.Chart().mark_bar(size=14, cornerRadiusTopRight=9, cornerRadiusBottomRight=9).encode(
    x=alt.X("charlson_early:Q", axis=alt.Axis(title="Charlson index")),
    x2=alt.X2("charlson_late:Q"),
    y=alt.Y("country:N", axis=alt.Axis(title="Country")),
    color=alt.Color("country:N", legend=alt.Legend(title="Country")),
    opacity=alt.Opacity("is_right_direction:Q", scale=triangle_opacity_scale, legend=None)
)

chart = alt.layer(left_bars, right_bars, data=df).facet(
    row=alt.Row("setting:N", header=alt.Header(title="Severity")),
)

chart = chart.properties(
    title={
        "text": "Charlson index from early wave to late wave",
        "subtitle": "Status: Pre-admission",
        "subtitleColor": "gray",
        "dx": 150
    }
)

chart = apply_theme(chart, title_anchor='start')

chart

# +
all_orig_df = orig_df.loc[orig_df["setting"] == "All"].copy()
all_orig_df["wave"] = all_orig_df["wave"].apply(lambda w: w[0].upper() + w[1:])

country_color_scale = alt.Scale(domain=["France", "Germany", "Italy", "Spain", "USA"], range=['#0072B2', '#E79F00', '#029F73', '#D45E00', '#CB7AA7'])

# Sort countries by their early wave index value, ascending

country_shift_sort = all_df.sort_values(by="charlson_early", ascending=True)["country"].tolist()

lines = alt.Chart().mark_line(point=True, size=3,).encode(
    y=alt.Y("charlson:Q", axis=alt.Axis(title="Charlson index")),
    x=alt.X("wave:N", axis=alt.Axis(title=None)),
    color=alt.Color("country:N", legend=alt.Legend(title="Country"), scale=country_color_scale)
)

chart = alt.layer(lines, data=all_orig_df).facet(
    column=alt.Column("country:N", header=alt.Header(title="Country", orient='bottom')),
)

chart = chart.properties(
    title={
        "text": "Charlson index from early wave to late wave",
        "subtitle": "Severity: All | Status: Pre-admission",
        "subtitleColor": "gray",
        "dx": 53,
    }
)

chart = apply_theme(chart, title_anchor='start')

chart

# +
all_orig_df = orig_df.loc[orig_df["setting"] == "All"].copy()
all_orig_df["wave"] = all_orig_df["wave"].apply(lambda w: w[0].upper() + w[1:])

country_color_scale = alt.Scale(domain=["France", "Germany", "Italy", "Spain", "USA"], range=['#0072B2', '#E79F00', '#029F73', '#D45E00', '#CB7AA7'])

# Sort countries by their early wave index value, ascending

country_shift_sort = all_df.sort_values(by="charlson_early", ascending=True)["country"].tolist()

lines = alt.Chart().mark_bar().encode(
    y=alt.Y("charlson:Q", axis=alt.Axis(title="Charlson index")),
    x=alt.X("wave:N", axis=alt.Axis(title=None)),
    color=alt.Color("country:N", legend=alt.Legend(title="Country"), scale=country_color_scale)
)

chart = alt.layer(lines, data=all_orig_df).facet(
    column=alt.Column("country:N", header=alt.Header(title="Country", orient='bottom')),
)

chart = chart.properties(
    title={
        "text": "Charlson index from early wave to late wave",
        "subtitle": "Severity: All | Status: Pre-admission",
        "subtitleColor": "gray",
        "dx": 53,
    }
)

chart = apply_theme(chart, title_anchor='start')

chart
# -







StarcoderdataJupyterScriptsDedupFiltered