Coverage for examples/simulation/csv_data_source.py: 59%

1import logging 

2from datetime import datetime, timedelta 

3 

4import matplotlib.pyplot as plt 

5import numpy as np 

6import pandas as pd 

7 

8import agentlib as ag 

9 

10 

11def create_sample_data(): 

12 date_today = datetime.now() 

13 time = pd.date_range(date_today, date_today + timedelta(minutes=10), freq="30S") 

14 data1 = np.sin(np.linspace(0, 2 * np.pi, len(time))) * 10 + 20 

15 data2 = np.cos(np.linspace(0, 2 * np.pi, len(time))) * 5 + 15 

16 df = pd.DataFrame({"timestamp": time, "temperature": data1, "humidity": data2}) 

17 df.set_index("timestamp", inplace=True) 

18 df.to_csv("sample_data.csv") 

19 return df 

20 

21 

22def run_example(log_level=logging.INFO, with_plots: bool = True, extrapolation: str = "constant"): 

23 # Set the log-level 

24 logging.basicConfig(level=log_level) 

25 

26 # Create sample data 

27 original_df = create_sample_data() 

28 

29 # Configure the agent 

30 agent_config = { 

31 "id": "DataSourceAgent", 

32 "modules": [ 

33 { 

34 "module_id": "DataSource", 

35 "type": "csv_data_source", 

36 "data": "sample_data.csv", 

37 "t_sample": 10, 

38 "outputs": [ 

39 {"name": "temperature", "shared": True}, 

40 {"name": "humidity", "shared": True}, 

41 ], 

42 "extrapolation": extrapolation, 

43 }, 

44 {"module_id": "Logger", "type": "AgentLogger"}, 

45 ], 

46 } 

47 

48 # Set up the environment and agent 

49 env_config = {"rt": False, "factor": 1, "t_sample": 1} 

50 env = ag.Environment(config=env_config) 

51 agent = ag.Agent(config=agent_config, env=env) 

52 

53 # Run the simulation 

54 env.run(until=600 * 2) # Run for 10 minutes (600 seconds) 

55 results = agent.get_results() 

56 

57 if with_plots: 

58 plot(results, original_df) 

59 

60 return {"DataSourceAgent": results} 

61 

62 

63def plot(results, original_df): 

64 # Get results from the AgentLogger 

65 logger_results = results["Logger"] 

66 # Convert original_df index to seconds from start 

67 original_df.index = (original_df.index - original_df.index[0]).total_seconds() 

68 # Plot the results 

69 fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 10)) 

70 # Temperature plot 

71 original_df["temperature"].plot( 

72 ax=ax1, 

73 label="Original Temperature", 

74 alpha=0.7, 

75 marker="o", 

76 ) 

77 logger_results["temperature"].plot( 

78 ax=ax1, 

79 label="DataSource Temperature", 

80 style="--", 

81 drawstyle="steps-post", 

82 ) 

83 ax1.set_title("Temperature Comparison") 

84 ax1.set_ylabel("Temperature (°C)") 

85 ax1.set_xlabel("Time (seconds)") 

86 ax1.legend() 

87 # Humidity plot 

88 original_df["humidity"].plot( 

89 ax=ax2, 

90 label="Original Humidity", 

91 alpha=0.7, 

92 marker="o", 

93 ) 

94 logger_results["humidity"].plot( 

95 ax=ax2, 

96 label="DataSource Humidity", 

97 style="--", 

98 drawstyle="steps-post", 

99 ) 

100 ax2.set_title("Humidity Comparison") 

101 ax2.set_ylabel("Humidity (%)") 

102 ax2.set_xlabel("Time (seconds)") 

103 ax2.legend() 

104 plt.tight_layout() 

105 plt.show() 

106 return results 

107 

108 

109if __name__ == "__main__": 

110 run_example(log_level=logging.INFO)