def main()

in prediction_generation/cpdbench_prophet.py [0:0]

• 39 lines of code
• 11 McCabe index (conditional complexity)

def main():
    args = parse_args()
    raw_args = copy.deepcopy(args)

    # Load the dataset (using a Python equivalent of your R helper function)
    data, mat = load_dataset(args.input)

    start_time = time.time()

    if args.Nmax == 'default':
        args.Nmax = 25
    else:
        args.Nmax = data['n_obs'] - 1  # Adjusted from 'original' to the main data structure

    # Check if 'series' is in data and extract appropriately
    if 'series' not in data or len(data['series']) == 0:
        exit_with_error(data, raw_args, vars(args), "No time series data available.")

    # Prepare the DataFrame for Prophet
    df = pd.DataFrame({
        'ds': data['time']['raw'],  # Time column
        'y': data['series'][0]['raw']  # Series values
    })

    # Handle logistic growth by adding 'cap'
    if args.growth == 'logistic':
        if args.cap is None:
            exit_with_error(data, raw_args, vars(args), "Capacity ('cap') must be provided for logistic growth.")
        
        # Add the capacity column to the DataFrame
        df['cap'] = args.cap

    # Fit the Prophet model
    try:
        model = Prophet(
            changepoint_range=args.ChangepointRange,
            n_changepoints=args.Nmax,
            weekly_seasonality=args.WeeklySeasonality,
            daily_seasonality=args.DailySeasonality,
            growth=args.growth,  # 'linear' or 'logistic'
            changepoint_prior_scale=args.ChangepointPriorScale,
            interval_width=args.IntervalWidth
        )
        model.fit(df)

        # Retrieve changepoints (timestamps)
        changepoints = model.changepoints
        locs = changepoints.index.tolist()

    except Exception as e:
        exit_with_error(data, raw_args, vars(args), str(e), __file__)

    stop_time = time.time()
    runtime = stop_time - start_time

    locs = convert_timestamps(locs)
    data = convert_timestamps(data)

    exit_success(data, raw_args, vars(args), locs, runtime, __file__)