Plotting with if-else conditional statements in python-CodePudding

I am trying to use go.scatter with my conditional statements.

A and df['T_orNonT'] are columns in my dataframe, df. If a row on "A" is less than or equal to 200, the column df['T_orNonT'] will show 'Non-T', otherwise it is 'T'

I want to plot them using go.scatter with 'T' or 'Non-T' showing up with different color. Here is my code:

import plotly.graph_objects as go
fig = go.Figure()
for i in range (0, length):
    if A[i] <=  200:
        df['T_or_NonT'].iloc[i] = 'Non-T'
        fig = go.Figure()
        fig.add_trace(go.Scatter(
             x = df['Date'],
             y = df['A'],
             mode ='markers',
             name='Non-T',
             marker=dict(color ='red')))
        fig.show()


    else:
        df['T_or_NonT'].iloc[i] = 'T'
        fig = go.Figure()
        fig.add_trace(go.Scatter(
             x = df['Date'],
             y = df['A'],
             mode ='markers',
             name='T',
             marker=dict(color ='green')))
        fig.show()

This should be the output:

  Date        A       T or Non-T
   07/21      201      T
   08/21      255      T
   09/21      198     Non-T

And then they will plot Date (monthly) vs Rainfall (which is the A column). The Ts are marked as red, and the Non-Ts are marked as green in the plot.

but I can't make it work. I want to know the right way to code this. by the way i am a python beginner-user.

PS. You can also suggest if there is another work-around

CodePudding user response：

There are many ways to do this, but I think the easiest is to have a column of colors for the decision results. The easiest way to do this is to have a color column of judgment results, and then draw a scatter plot with the data extracted by the judgment condition items for that data frame.

import pandas as pd
import numpy as np
import plotly.graph_objects as go

df = pd.DataFrame({'Date': pd.date_range('2018-01-21','2021-01-21',freq='1m'),
                  'A': np.random.randint(150,250, 36)})

df['T_or_NonT'] = np.where(df['A'] >= 200,'T','Non-T')
df['color'] = np.where(df['A'] >= 200,'red','green')

fig = go.Figure()

for t,c in zip(df['T_or_NonT'].unique(), df['color'].unique()):
    dfs = df[df['T_or_NonT'] == t]
    fig.add_trace(go.Scatter(
        x = dfs['Date'],
        y = dfs['A'],
        mode = 'markers',
        name = t,
        marker = dict(
            color = c
        )
    ))
    
fig.show()