dancing-task-modeling/plot_dancing_task.m at main · sevenbanisch/dancing-task-modeling · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
function plot_dancing_task(Obs_d, dyad, env) % version for only one dynamic agent
    figure(1)
    clf

    % -------------------------------
    % Main plot
    % -------------------------------
    ax1 = axes('Position', [0.10 0.11 0.68 0.815]);

    plot(ax1, Obs_d, 'LineWidth', 2, 'Color', dyad.A.color)
    ylim(ax1, [0 env.dmax])
    yticks(ax1, 0:1:env.dmax)
    hold(ax1, 'on')
    yline(ax1, dyad.A.delta, 'r--', 'Preferred distance ('+(dyad.A.name)+')', 'LineWidth', 2)
    yline(ax1, dyad.B.delta, 'r--', 'Preferred distance ('+(dyad.B.name)+')', 'LineWidth', 2, 'LabelVerticalAlignment', 'bottom')
    hold(ax1, 'off')

    legend(ax1, {'Distance'}, 'Location', 'best')
    set(ax1, 'FontSize', 14)

    % -------------------------------
    % Calculate rewards
    % -------------------------------
    y = linspace(0, env.dmax, 400).';
    reward = arrayfun(@(yy) preference(0, yy, dyad.A.delta, dyad.A.deltarange), y);
    reward = reward(:);
    reward(~isfinite(reward)) = 0;

    % -------------------------------
    % Right axis: small color bar
    % -------------------------------
    ax2 = axes('Position', [0.82 0.11 0.03 0.815]);

    imagesc(ax2, [0 1], [y(1) y(end)], reward)
    set(ax2, 'YDir', 'normal')
    ylim(ax2, [0 env.dmax])
    xlim(ax2, [0 1])

    % Disable labels on these axes
    xticks(ax2, [])
    yticks(ax2, [])
    set(ax2, 'Box', 'on', 'FontSize', 12)

    % -------------------------------
    % Colormap
    % -------------------------------
    n = 256;
    half = floor(n/2);

    red_to_yellow = [ ...
        ones(half,1), ...
        linspace(0,1,half)', ...
        zeros(half,1)];

    yellow_to_green = [ ...
        linspace(1,0,n-half)', ...
        ones(n-half,1), ...
        zeros(n-half,1)];

    cmap = [red_to_yellow; yellow_to_green];
    colormap(ax2, cmap)

    cmax = max(abs(reward));
    if cmax == 0
        cmax = 1;
    end
    clim(ax2, [-cmax cmax])

    % -------------------------------
    % Labels right
    % -------------------------------
    labelY = 0:1:env.dmax;
    labelReward = arrayfun(@(yy) preference(0, yy, dyad.A.delta, dyad.A.deltarange), labelY);

    for i = 1:numel(labelY)
        text(ax2, 1.15, labelY(i), sprintf('%.2f', labelReward(i)), ...
            'VerticalAlignment', 'middle', ...
            'HorizontalAlignment', 'left', ...
            'FontSize', 11, ...
            'Clipping', 'off');
    end

    text(ax2, 3.6, env.dmax/2, 'Reward', ...
        'Rotation', 90, ...
        'HorizontalAlignment', 'center', ...
        'VerticalAlignment', 'middle', ...
        'FontSize', 13, ...
        'Clipping', 'off');

    linkaxes([ax1 ax2], 'y')
    set(gcf, 'WindowStyle', 'docked')

end