1つ目のデータと、それ以降のデータの算出計算が異なる場合の計算式の書き方

Question

前提・実現したいこと

「平均足」という指標の計算をしたい。
（最初（厳密には2番め）のデータと、それ以降のデータとで、計算式が異なる場合の計算を、どのように行うか）

困っていること、及び該当のソースコード

以下のようなデータがあった場合、

#      date      open   high   low    close
#0  2020-01-14   1810   1849   1785   1807
#1  2020-01-15   1800   1832   1776   1805
#2  2020-01-16   1830   1857   1805   1833
#3  2020-01-17   1864   1943   1813   1915

求めたい値を
・平均足の終値(HA_close)
・平均足の始値(HA_open)
とすると、求めたい計算式は以下のようになります。

平均足の終値(HA_close) = (当日始値+当日高値+当日安値+当日安値) / 4
→ df['HA_close'] = (df['open'] + df['high'] + df['low'] + df['close']) / 4

平均足の始値(HA_open)
【1日目の場合】※上記ケースだと2020-01-15
= (前日始値+前日高値+前日安値+前日終値) / 4
→ df['HA_open'] = (df['open'].shift(1) + df['high'].shift(1) + df['low'].shift(1) + df['close'].shift(1)) / 4

【2日目以降の場合】
= (前日の平均足の始値+前日の平均足の終値) / 2
→ (df['HA_open'].shift(1) + df['HA_close'].shift(1)) / 2

平均足の終値(HA_close)の方は単純な計算なので問題ないのですが、
平均足の始値(HA_open)の方は、最初（厳密には2番め）のデータと、それ以降のデータとで、計算式が異なるので、どのようにコードを書くのがいいか頭を悩ませています。

#      date      open   high   low    close   HA_open   HA_close
#0  2020-01-14   1810   1849   1785   1807               1813
#1  2020-01-15   1800   1832   1776   1805     1813      1803
#2  2020-01-16   1830   1857   1805   1833     1808      1831
#3  2020-01-17   1864   1943   1813   1915     1820      1884

最終的には、上記のような形で出力させたく、以下のような形で書いたものの、HA_openは表示されませんでした。

for n in range(0, len(df)):
    if n == 0:
        df['HA_open'].iloc[n,:] = 0
    elif n == 1:
        df['HA_open'].iloc[n,:] = (df['open'].shift(1) + df['high'].shift(1) + df['low'].shift(1) + df['close'].shift(1)) / 4
    else:
        df['HA_open'] = (df['HA_open'].shift(1) + df['HA_close'].shift(1)) / 2

まだ知識が浅く、理解が間違っていたり、不足していたり、応用が効かない状態のため、質問の仕方自体もままならず、うまく伝わっていないかもしれませんが、アドバイスいただけると助かります。お手数おかけしますが、よろしくお願いします。

Answer 1

pandas の動作としては、

• 先頭行で Shift(1)した時の値：NaN
• NaN 同士の値の平均：NaN
• NaN を含めた数列の平均：NaN値を除いた数での平均

なので、そもそも条件文が不要ですし、１行目や２行目とそれ以外を分ける必要すらないのではないでしょうか。

• １日目：前の行がないので全ての'前日の平均足の始値'、'前日の平均足の終値'は共にNaNとなるので、その平均値である'当日 の平均足の始値'はNaNと計算される
• ２日目：上記の計算より'前日の平均足の始値'はNaNとなるので、'前日の平均足の終値'との平均値である'当日の平均足の始値'はそのまま'前日の平均足の終値'と計算される

まとめるとこんな感じです。

import pandas as pd
import numpy as np

# 初期データ生成
df = pd.DataFrame(
    [['2020-01-14', 1810, 1849, 1785, 1807],
     ['2020-01-15', 1800, 1832, 1776, 1805],
     ['2020-01-16', 1830, 1857, 1805, 1833],
     ['2020-01-17', 1864, 1943, 1813, 1915]],
    columns=['date','open','high','low','close'])
df = df.set_index('date')

# ここから平均足の始値と平均足の終値の計算
df['HA_open'] = np.nan
df['HA_close'] = df.mean(axis=1)
for idx, row in df.iterrows():
    df.loc[idx, 'HA_open'] = df.shift(1).loc[idx, ['HA_open', 'HA_close']].mean()

print(df)
#            open  high   low  close   HA_open  HA_close
#date
#2020-01-14  1810  1849  1785   1807       NaN   1812.75
#2020-01-15  1800  1832  1776   1805  1812.750   1803.25
#2020-01-16  1830  1857  1805   1833  1808.000   1831.25
#2020-01-17  1864  1943  1813   1915  1819.625   1883.75

Comments

[qsuke]

詳細にありがとうございます！
上記で無事解決しました。

ご指摘いただいたように、難しく考えすぎていたように思います。。とてもすっきりしました。

Answer 2

もう既に回答が出ているようですが、せっかく処理を作ってみたので掲載してみます。

好みで選んでもよいのではと思います。

# 初期データ生成
n_array = [
    ['2020-01-14', 1810, 1849, 1785, 1807],
    ['2020-01-15', 1800, 1832, 1776, 1805],
    ['2020-01-16', 1830, 1857, 1805, 1833],
    ['2020-01-17', 1864, 1943, 1813, 1915]
]
df = pd.DataFrame(n_array, columns=['date','open','high','low','close'])
df['HA_open'] = None
df['HA_close'] = (df['open'] + df['high'] + df['low'] + df['close']) / 4 # 入れるの忘れてた

# HA_open HA_close 計算
df_len = len(df)
for n in range(0, df_len):
    if n == 0:
        HA_open = 0
    elif n == 1:
        HA_open = (df.loc[n - 1][['open', 'high', 'low', 'close']].sum()) / 4
    else:
        HA_open = (df.loc[n - 1][['HA_open', 'HA_close']].sum()) / 2

    df.loc[n, 'HA_open'] = HA_open

df

Comments

[siruku6]

ソースコードに一部書き忘れがあったので後から書き足しています...

Answer 3

n=0,n=1の場合はループの中で処理する必要がなさそうですよね。
中の数式についてちゃんと確認しておらず動作チェックもしてませんが、雰囲気として次のような感じにできるのでないでしょうか。

df['HA_open'][0] = 0
df['HA_open'][1] = (df['open'].shift(1) + df['high'].shift(1) + df['low'].shift(1) + df['close'].shift(1)) / 4
for n in range(2, len(df)):
    df['HA_open'][n] = (df['HA_open'].shift(1) + df['HA_close'].shift(1)) / 2

なお、本件では必要なさそうに見えますが、場合分けを含む式を立てたい場合は三項演算子というのが使えます。
https://note.nkmk.me/python-if-conditional-expressions/

Comments

[qsuke]

さっそくご指摘いただきありがとうございます！
結構シンプルに行けるんですね…汗

さっそくアドバイスいただいた方法でやってみたところ、
KeyError: 'HA_open'
というエラーが出て、それ以降も
KeyError: "['HA_open'] not in index"　
となってしまい、うまく処理しきれていませんでした。。

この計算部分以外のところで問題があるのでしょうか？

df['HA_close']は問題なく計算できているので、df['HA_open']の計算部分だと思うのですが、何かこのエラーから気になる点はありますでしょうか？

[KojiDoi]

dfにまだ'HA_close’列がないのですね。その場合は、df.assign()などで追加する必要があるでしょう。
もしくは、HA_closeに入れるべきデータを先に一次元配列として作り、
df['HA_close']=一次元配列の名前
のようにすればいけるのではないかと思います。

[qsuke]

ご丁寧にありがとうございます！
ただ、いろいろ調べてやってみているのですが、アドバイスいただいた内容をきちんと理解できないぐらいのレベルでして、どうにもうまくいかない状態です…

かなり初歩的なレベルのことだと思いますので、今一度、基本的内容に立ち戻って、勉強しながらトライしてみたいと思います。ありがとうございました。